#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   splice_main.py
@Author  :   Liu Kaiyu
'''

import os
from PyQt5 import QtWidgets,QtCore,QtGui
from splice import Ui_Splice
from three_dialogs import HelpMenu,PreSet
from generators import FFMPEGGenerator,OverlayCMDGenerator,SpliceCMDGenerator
from splice_thread import SpliceThread
from myfilebrowser_main import MyFileDialog

# This will be shown on UI
date_modified = '30/08/2019'

# Hard code some commonly used image extensions
still_img_extensions = ['tif','tiff','png','jpeg','gif','exr','tga','jpg','jif']

################################################################################################################################################
# Customized Erros
class MyError(Exception):
    def __init__(self,message):
        self.message = message

################################################################################################################################################
# Define a global method, this will be used in FileList class below
def split_folder_file(file_path):
    folder,file_name_with_extension = os.path.split(file_path)
    file_name,extension = os.path.splitext(file_name_with_extension)

    return folder,file_name,extension

################################################################################################################################################
################################################################################################################################################
################################################################################################################################################
class FileList(QtWidgets.QTableWidget):

    refresh_signal = QtCore.pyqtSignal()

    def __init__(self,parent):
        super(FileList,self).__init__(parent)
        self.setAcceptDrops(True)

        # Some Settings, copy from UI codes
        # Default has one row, and two columns ("file name" and "duration" )
        self.setGeometry(QtCore.QRect(10, 50, 431, 141))
        self.setColumnCount(2)
        self.setRowCount(1)
        self.resizeRowsToContents()
        self.setHorizontalHeaderLabels(['File Name','Duration'])
        self.setColumnWidth(0,300)
        self.verticalHeader().setDefaultSectionSize(15)
        self.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.setGridStyle(QtCore.Qt.NoPen)
        self.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)

    ##############################################################################################################
    def dragMoveEvent(self, event):
        if event.mimeData().hasUrls:
            event.setDropAction(QtCore.Qt.CopyAction)
            event.accept()
        else:
            event.ignore()

    ##############################################################################################################
    def dragEnterEvent(self, e):
        m = e.mimeData()
        # File path is considered as URL
        if m.hasUrls():
            e.accept()
        else:
            e.ignore()

    ##############################################################################################################
    def dropEvent(self, e):
        m = e.mimeData()
        if m.hasUrls():
            for file in m.urls():
                file_path = file.toLocalFile()

                file_name = split_folder_file(file_path)[1]
                extension = split_folder_file(file_path)[2]

                self.parent().file_list_dd[file_name + extension] = {'file_path':'%s'%file_path}
                self.parent().get_file_identity(file_path,file_name + extension)
                self.parent().check_valid_image_sequence(file_name + extension)

            self.refresh_signal.emit()
            print('refresh signal emitted')

################################################################################################################################################
################################################################################################################################################
################################################################################################################################################

class SpliceTool(QtWidgets.QMainWindow,Ui_Splice):
    def __init__(self):
        QtWidgets.QMainWindow.__init__(self)
        self.setupUi(self)

        # Re-write drop button
        self.file_list = FileList(self)
        self.file_list.refresh_signal.connect(self.refresh_file_list)

        # Update date label
        global date_modified
        _translate = QtCore.QCoreApplication.translate
        self.update_label.setText(_translate("MainWindow", "Update: %s"%date_modified))
        self.setWindowIcon(QtGui.QIcon(r'.\resources\milo.png'))

        # Connections
        self.input_browse.clicked.connect(lambda: self.browse_path(inp = True))
        self.output_browse.clicked.connect(lambda: self.browse_path(out = True))
        self.clear_btn.clicked.connect(self.clear_file_list_box)
        self.up_btn.clicked.connect(lambda: self.move_file('up'))
        self.down_btn.clicked.connect(lambda: self.move_file('down'))
        self.close_btn.clicked.connect(lambda: self.close())
        self.delete_btn.clicked.connect(self.delete_item)
        self.splice_btn.clicked.connect(self.splice_btn_function)
        self.overlay_checkbox.stateChanged.connect(self.use_overlay)
        self.file_list.itemClicked.connect(self.check_is_still_img)
        self.ok_btn.clicked.connect(self.ok_btn_function)
        self.clear_all_btn.clicked.connect(self.clear_all)

        # Store files and cmds
        self.file_list_dd = {}
        self.cmd_list = []

        # Initialize tool button menu and setting
        menu = QtWidgets.QMenu(self)
        self.help_Act = QtWidgets.QAction('Help',self)
        self.pre_setting_Act = QtWidgets.QAction('Pre Sets',self)

        menu.addAction(self.help_Act)
        menu.addAction(self.pre_setting_Act)
        self.toolButton.setMenu(menu)

        # Connect dialogs to mainwindows
        self.help_menu_dialog = HelpMenu()
        self.help_Act.triggered.connect(lambda: self.help_menu_dialog.show())

        self.preset_dialog = PreSet()
        self.pre_setting_Act.triggered.connect(lambda: self.preset_dialog.show())

        self.PRESETS = self.preset_dialog.read_preset_book()
        self.refresh_preset_dropdown(self.PRESETS)
        self.preset_dialog.refresh_signal.connect(lambda:  self.refresh_preset_dropdown(self.preset_dialog.read_preset_book()))

        self.fmt_generator = FFMPEGGenerator()
        self.splice_cmd_generator = SpliceCMDGenerator()

        # Thread
        self.threadpool = QtCore.QThreadPool()

    ################################################################################################################################################
    def create_a_thread(self):
        self.splice_tool = SpliceThread()
        self.splice_tool.signals.finished.connect(self.finish)

    ################################################################################################################################################
    def browse_path(self,inp = False, out = False):
        path = ''

        supported_video_fmt = self.fmt_generator.get_muxing_supported_fmt()

        formats = ''
        for each_fmt in supported_video_fmt:
            formats += '%s(*.%s);;'%(each_fmt.strip(),each_fmt.strip())

        if inp :
            file_dialog = MyFileDialog()
            path_list = file_dialog.exec_()

            if path_list:
                for path in path_list:
                    print(path)
                    self.add_files_to_list(path)

        elif out:
                path = QtWidgets.QFileDialog.getSaveFileName(self,'Choose File','.','%s'%formats)[0]

                if path != '':
                    path = path.replace('/','\\')
                    self.output_dropdown.clear()
                    self.output_dropdown.addItem(path)

    ################################################################################################################################################
    # Add files to file list box, and store the file path to 'file_list_dd'
    def add_files_to_list(self,file_path):
        file_name = os.path.split(file_path)[1].strip()
        print(file_name)

        self.file_list_dd[file_name] = {'file_path':'%s'%file_path}
        self.get_file_identity(file_path,file_name)

        # Check valid image sequence
        self.check_valid_image_sequence(file_name)

        self.refresh_file_list()

    ################################################################################################################################################
    def check_valid_image_sequence(self,file_name):

        if self.file_list_dd[file_name]['identity'] == 'image_sequence':
            folder_path = self.file_list_dd[file_name]['file_path']
            name_set = set()
            for each_img in os.listdir(folder_path):
                img_without_extension = os.path.splitext(each_img)[0]

                i = 1
                temp = []
                while True:

                    if i == len(img_without_extension):
                        break

                    if img_without_extension[-i].isdigit():
                        i += 1
                        continue
                    else:
                        temp.insert(0,img_without_extension[:-i])
                        break

                image_sequence_name = ''.join(i for i in temp)
                name_set.add(image_sequence_name)

            if len(name_set) > 1:
                QtWidgets.QMessageBox.about(self,' Warning ','%s has more than one image sequence'%file_name )
                del self.file_list_dd[file_name]

    ################################################################################################################################################
    def get_file_name_from_list(self,row_content):
        if '|' in row_content:
            file_name = row_content.split('|')[1].strip()
        else:
            file_name = row_content.strip()

        return file_name

    ################################################################################################################################################
    def refresh_file_list(self):
        print('refresh signal received')

        self.file_list.setRowCount(len(self.file_list_dd.keys()))

        for index,each_file in enumerate(self.file_list_dd.keys()):
            file_name = QtWidgets.QTableWidgetItem(each_file)
            self.file_list.setItem(index,0,file_name)            #Format : (row,col,item)

            if self.file_list_dd[each_file]['identity'] == 'still_image':
                if 'duration' in self.file_list_dd[each_file]:
                    duration = QtWidgets.QTableWidgetItem(self.file_list_dd[each_file]['duration'])
                else:
                    duration = QtWidgets.QTableWidgetItem('!')
                self.file_list.setItem(index,1,duration)

            else:
                NA = QtWidgets.QTableWidgetItem('NA')
                self.file_list.setItem(index,1,NA)

    ################################################################################################################################################
    def clear_file_list_box(self):
        self.file_list.clearContents()
        self.file_list.setRowCount(1)
        self.duration_inp.clear()
        self.file_list_dd ={}
        self.cmd_list = []

    ################################################################################################################################################
    def clear_all(self):
        self.clear_file_list_box()
        self.output_dropdown.clear()
        self.progressBar.setValue(0)

    #############################################################################################################
    def use_overlay(self):
        if self.overlay_checkbox.isChecked():
            self.overlay_group.setEnabled(True)
        else:
            self.overlay_group.setEnabled(False)

    ################################################################################################################################################
    def refresh_preset_dropdown(self,PRESETS):
        self.preset_dropdown.clear()
        self.preset_dropdown.addItem('None')

        for each_set in sorted(PRESETS.keys()):
            self.preset_dropdown.addItem(each_set)
        print(self.PRESETS)

    ################################################################################################################################################
    def finish(self,signal):
        if signal == 'done':
            self.progressBar.setRange(0,1)
            self.progressBar.setValue(1)
            QtWidgets.QMessageBox.about(self,' ','Done !')
            self.prevent_misclicking(finish=True)

    #############################################################################################################
    def closeEvent(self,event):
        '''
        overwrite closeEvent, pop a comfirmation window before closing
        '''
        reply = QtWidgets.QMessageBox.question(self,'Warning','Are you sure to Exit',QtWidgets.QMessageBox.Yes|QtWidgets.QMessageBox.No)
        if reply == QtWidgets.QMessageBox.Yes:
            self.help_menu_dialog.close()
            self.preset_dialog.close()
            self.close()
        else:
            event.ignore()

    #############################################################################################################
    def move_file(self,direction):
        selected_row = self.file_list.currentRow()

        # If no row has been selected, return value is -1
        if selected_row == -1 :
            QtWidgets.QMessageBox.about(self,' ','Please select a file !')
        else:
            file_name = self.file_list.item(selected_row,0).text()
            duration = self.file_list.item(selected_row,1).text()

            if direction == 'up' and selected_row != 0:
                self.file_list.insertRow(selected_row-1)
                self.file_list.setItem(selected_row-1,0,QtWidgets.QTableWidgetItem(file_name))
                self.file_list.setItem(selected_row-1,1,QtWidgets.QTableWidgetItem(duration))
                self.file_list.removeRow(selected_row+1)
                self.file_list.setCurrentCell(selected_row-1,0)

            elif direction == 'down' and selected_row != len(self.file_list_dd.keys())-1:
                self.file_list.insertRow(selected_row+2)
                self.file_list.setItem(selected_row+2,0,QtWidgets.QTableWidgetItem(file_name))
                self.file_list.setItem(selected_row+2,1,QtWidgets.QTableWidgetItem(duration))
                self.file_list.removeRow(selected_row)
                self.file_list.setCurrentCell(selected_row+1,0)

    ################################################################################################################################################
    def delete_item(self):
        selected_row = self.file_list.currentRow()

        if selected_row == -1 :
            QtWidgets.QMessageBox.about(self,' ','Please select a item !')

        else:
            file_name = self.file_list.item(selected_row,0).text()

            reply = QtWidgets.QMessageBox.question(self,'Warning','Are you sure to delete [%s]?'%file_name,QtWidgets.QMessageBox.Yes|QtWidgets.QMessageBox.No)

            if reply == QtWidgets.QMessageBox.Yes:
                del self.file_list_dd[file_name]
                self.duration_inp.clear()

                self.file_list.removeRow(selected_row)
                print(self.file_list_dd)

    ################################################################################################################################################
    def get_file_identity(self,file_path,file_name):
        global still_img_extensions
        if os.path.isfile(file_path):
            if file_name.split('.')[-1].strip().lower() in still_img_extensions:
                identity = 'still_image'
            else:
                identity = 'video'
        elif os.path.isdir(file_path):
            identity = 'image_sequence'
        else:
            QtWidgets.QMessageBox.about(self,' ',' Unknown File ')

        print('%s is %s'%(file_path,identity))

        self.file_list_dd[file_name]['identity'] = identity
        print(self.file_list_dd)

    ################################################################################################################################################
    # Check if a file is a still image
    def check_is_still_img(self):
        global still_img_extensions

        self.duration_inp.clear()
        selected_row = self.file_list.currentRow()
        file_name = self.file_list.item(selected_row,0).text()

        if file_name.split('.')[-1].strip().lower() in still_img_extensions:

            self.duration_inp.setEnabled(True)
            self.ok_btn.setEnabled(True)

            if 'duration' in self.file_list_dd[file_name].keys():
                self.duration_inp.setText(self.file_list_dd[file_name]['duration'])

        else:

            self.duration_inp.setEnabled(False)
            self.ok_btn.setEnabled(False)

    ################################################################################################################################################
    def ok_btn_function(self):
        if not self.duration_inp.text().strip().isdigit():
            QtWidgets.QMessageBox.about(self,' ERROR ',' Please input a valid duration ')
        elif self.duration_inp.text().strip() == '0':
            QtWidgets.QMessageBox.about(self,' ERROR ',' Duration must be larger than zero ')
        else:
            selected_row = self.file_list.currentRow()
            if selected_row == -1:
                QtWidgets.QMessageBox.about(self,' ERROR ',' Please select a file ')
            else:
                file_name = self.file_list.item(selected_row,0).text()
                self.file_list_dd[file_name]['duration']= self.duration_inp.text().strip()
                for obj in self.file_list.findItems(file_name,QtCore.Qt.MatchExactly):
                    self.file_list.setItem(obj.row(),1,QtWidgets.QTableWidgetItem(self.duration_inp.text().strip()))

                self.duration_inp.clear()

            print(self.file_list_dd)

            #self.refresh_file_list()

    ################################################################################################################################################
    def get_resolution(self):
        w = self.w_inp.text().strip()
        h = self.h_inp.text().strip()
        resolution = w + 'x' + h

        return resolution

    ################################################################################################################################################
    def splice_btn_function(self):
        if self.check_errors():

            self.create_a_thread()

            # Create a sequence list
            num_of_files = self.file_list.rowCount()
            sequence_list = []
            for i in range(num_of_files):
                file_name = self.file_list.item(i,0).text()
                sequence_list.append(file_name)
            print('file list in GUI is : \n %s'%sequence_list)

            # Get selcected preset parameters
            selected_preset = self.preset_dropdown.currentText()
            parameter_list = self.PRESETS[selected_preset].split(',')

            # Pass values to back end thread

            if self.overlay_checkbox.isChecked():
                print('overlay is checked')
                self.splice_tool.overlay_flag = True

                project_name = self.project_inp.text().strip()
                author = self.author_inp.text().strip()
                task = self.task_inp.text().strip()
                version = self.version_inp.text().strip()

                self.splice_tool.overlay_info_list = [project_name,author,task,version]

            self.splice_tool.preset_parameter_list = parameter_list
            self.splice_tool.output_file = self.output_dropdown.currentText().strip()
            self.splice_tool.target_resolution = self.get_resolution()
            self.splice_tool.sequence_list = sequence_list
            self.splice_tool.file_list_dd = self.file_list_dd

            self.progressBar.setRange(0,0)
            self.threadpool.start(self.splice_tool)
            self.prevent_misclicking()

    ################################################################################################################################################
    # Disable some buttons when it is in splicing process
    def prevent_misclicking(self,finish=False):
        self.output_group.setEnabled(finish)
        self.input_browse.setEnabled(finish)
        self.preser_group.setEnabled(finish)
        self.file_list.setEnabled(finish)
        self.clear_btn.setEnabled(finish)
        self.delete_btn.setEnabled(finish)
        self.duration_inp.setEnabled(finish)
        self.ok_btn.setEnabled(finish)
        self.overlay_checkbox.setEnabled(finish)
        self.splice_btn.setEnabled(finish)
        self.overlay_group.setEnabled(finish)
        self.clear_all_btn.setEnabled(finish)
        self.use_overlay()

    ################################################################################################################################################
    def check_errors(self):

        try:
            num_of_files = len(self.file_list_dd.keys())
            print('%s files in file list'%num_of_files)
            if not num_of_files:
                raise MyError(' File list is empty ! ')

        except MyError as e:
            QtWidgets.QMessageBox.about(self,' Warning ',e.message)

        else:
            try:

                if not self.output_dropdown.currentText().strip():
                    raise MyError(' Please specify output file ')

                if self.preset_dropdown.currentText().strip() == 'None':
                    raise MyError(' Please select a Pre-Set ')

                if not self.w_inp.text().strip().isdigit():
                    raise MyError(' Please input a valid resolution Width')

                if not self.h_inp.text().strip().isdigit():
                    raise MyError(' Please input a valid resolution Height')

                for i in range(num_of_files):
                    file_name = self.file_list.item(i,0).text()
                    # Identity is at index 1 {file name:[path,identity]}
                    identity = self.file_list_dd[file_name]['identity']

                    # Specific duration must be assigned to still images
                    if identity == 'still_image':
                        # If still image is not assigned with a duration
                        if 'duration' not in self.file_list_dd[file_name].keys():
                            raise MyError(' Please assign a duration for File %s '%file_name)

                        # Check if duration is a digit
                        elif not self.file_list_dd[file_name]['duration'].isdigit():
                            raise MyError(' Still image %s has invalid duration '%file_name)

                    else:
                        continue

            except MyError as e:
                print(e.message)
                QtWidgets.QMessageBox.about(self,' Warning ',e.message)

            else:
                return True

if __name__ == '__main__':
    os.environ["QT_AUTO_SCREEN_SCALE_FACTOR"] = "1"
    app = QtWidgets.QApplication(os.sys.argv)
    app.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)

    splice_tool = SpliceTool()
    splice_tool.show()
    app.exec_()


from typing import Optional, Dict

from allennlp.data import Vocabulary
from allennlp.modules import FeedForward, Seq2SeqEncoder, TextFieldEmbedder
from allennlp.modules import InputVariationalDropout, TimeDistributed
from allennlp.modules.attention import DotProductAttention
from allennlp.models.model import Model
from allennlp.nn import InitializerApplicator, RegularizerApplicator
from allennlp.nn import util
from allennlp.training.metrics import CategoricalAccuracy, F1Measure
import numpy as np
import torch
import torch.nn.functional as F
from torch.nn import Linear, Dropout
from torch.nn.parameter import Parameter

@Model.register("attention_dialect_classifier")
class AttentionDialectClassifier(Model):
    def __init__(self,
                 vocab: Vocabulary,
                 text_field_embedder: TextFieldEmbedder,
                 text_encoder: Seq2SeqEncoder,
                 classifier_feedforward: Optional[FeedForward] = None,
                 dropout: Optional[float] = 0.0,
                 code_switching_regularizer: Optional[float] = 0.0,
                 bivalency_regularizer: Optional[float] = 0.0,
                 initializer: InitializerApplicator = InitializerApplicator(),
                 regularizer: Optional[RegularizerApplicator] = None) -> None:
        '''
        :param dropout: The amount of dropout to apply. Dropout is applied
                        after each non-linear layer and the word embeddings
                        lookup. Two types of dropout are applied, variational
                        dropout is applied if the input is to the dropout is
                        a sequence of vectors (each vector in the sequence
                        representing a word), and normal dropout if the input
                        is a vector.
        :param code_switching_regularizer: The weight associated to the code
                                           switching lexicon regulisation the
                                           lower the less affect it has. This
                                           requires that the dataset reader is
                                           going to supply the code switching
                                           arrays for the forward function of
                                           this class. If set a good values is
                                           0.001
        :param bivalency_regularizer: The weight associated to the bivalency
                                      regulisation the lower the less affect it
                                      has. This requires that the dataset
                                      reader is going to supply the bivalency
                                      arrays for the forward function of this
                                      class.
        '''
        super().__init__(vocab, regularizer)
        self._naive_dropout = Dropout(dropout)
        self._variational_dropout = InputVariationalDropout(dropout)

        self.text_field_embedder = text_field_embedder
        self.num_classes = self.vocab.get_vocab_size("labels")
        self.text_encoder = text_encoder
        text_encoder_dim = text_encoder.get_output_dim()
        # Attention parameters
        self.project_encoded_text = TimeDistributed(Linear(text_encoder_dim,
                                                           text_encoder_dim))
        self.attention_vector = Parameter(torch.Tensor(text_encoder_dim))
        self.reset_parameters()
        self.attention_layer = DotProductAttention(normalize=True)

        self.classifier_feedforward = classifier_feedforward
        output_dim = text_encoder_dim
        if classifier_feedforward:
            output_dim = classifier_feedforward.get_output_dim()
        self.label_projection = Linear(output_dim, self.num_classes)
        self.metrics = {
                "accuracy": CategoricalAccuracy()
        }
        self.code_switching_regularizer = code_switching_regularizer
        self.bivalency_regularizer = bivalency_regularizer
        self.loss = torch.nn.CrossEntropyLoss()
        initializer(self)

    def reset_parameters(self):
        '''
        Intitalises the attnention vector
        '''
        torch.nn.init.uniform_(self.attention_vector, -0.01, 0.01)

    def lexicon_loss(self, lexicon_array: torch.Tensor,
                     attention_weights: torch.Tensor,
                     text_mask: torch.Tensor) -> float:
        '''
        :param lexicon_array: Tensor of shape [batch, sequence_length] that
                              contains either 0, -1, or 1 values where -1
                              has to be for the whole sample in the batch
                              which would indicate not to perform attention
                              regulisation else 0 and 1 indicate important words
                              in the sequence that should be attended to.
        :param attention_weights: Tensor of shape [batch, sequence_length]
                                  that contains weights associating to
                                  the importance of the words in the sequence.
        :param text_mask: Tensor of shape [batch, sequence_length]. Mask
                          denoting whether there is a word or not at that
                          position.
        :returns: The difference between the lexicon_array denotating words
                  that should important to the network and those that the
                  network think are important via it's attention mechansim.
                  The difference is computed via cross entropy.
        '''
        # Mask is required to not add loss when the label is MSA
        lexicon_mask = lexicon_array.sum(1)
        lexicon_mask = (lexicon_mask >= 0).float()
        # Calculating the lexicon loss
        lexicon_attention_weights = attention_weights.squeeze()
        lexicon_softmax = util.masked_softmax(lexicon_array, text_mask)
        # This is to stop NANS from performing log on 0.
        lexicon_attention_weights = lexicon_attention_weights + 1e-8
        # Perform cross entropy like in the paper to get a loss function
        # that represents disagreement between the two vectors
        lexicon_loss = lexicon_softmax * torch.log(lexicon_attention_weights)
        lexicon_loss = lexicon_loss * text_mask.float()
        lexicon_loss = lexicon_loss.abs()
        lexicon_loss = lexicon_loss.sum(1)

        lexicon_loss = lexicon_loss * lexicon_mask
        lexicon_loss = lexicon_loss.sum()
        return lexicon_loss

    def forward(self,
                text: Dict[str, torch.LongTensor],
                label: torch.LongTensor = None,
                code_switching_array: torch.Tensor = None,
                bivalency_array: torch.Tensor = None
                ) -> Dict[str, torch.Tensor]:
        # Embed text
        embedded_text = self.text_field_embedder(text)
        embedded_text = self._variational_dropout(embedded_text)
        text_mask = util.get_text_field_mask(text)
        # Encode text
        encoded_text = self.text_encoder(embedded_text, text_mask)
        encoded_text = self._variational_dropout(encoded_text)
        #
        # ATTENTION
        #
        # Project each hidden timestep into a new hidden space
        # before attention
        projected_encoding = self.project_encoded_text(encoded_text)
        projected_encoding = torch.tanh(projected_encoding)
        projected_encoding = self._variational_dropout(projected_encoding)
        # Apply attention over the project encoded text
        batch_size = text_mask.shape[0]
        attention_vector = self.attention_vector.unsqueeze(0).expand(batch_size, -1)
        attention_weights = self.attention_layer(attention_vector,
                                                 projected_encoding, text_mask)
        attention_weights = attention_weights.unsqueeze(-1)
        weighted_encoded_text = attention_weights * encoded_text
        weighted_encoded_vec = weighted_encoded_text.sum(1)
        weighted_encoded_vec = self._naive_dropout(weighted_encoded_vec)

        if self.classifier_feedforward:
            weighted_encoded_vec = self.classifier_feedforward(weighted_encoded_vec)
        logits = self.label_projection(weighted_encoded_vec)
        class_probabilities = F.softmax(logits, dim=-1)

        output_dict = {"class_probabilities": class_probabilities}

        if label is not None:

            loss = self.loss(logits, label)

            if code_switching_array is not None and self.training:
                code_switch_loss = self.lexicon_loss(code_switching_array,
                                                     attention_weights, text_mask)
                loss = loss + (self.code_switching_regularizer * code_switch_loss)
            if bivalency_array is not None and self.training:
                bivalency_loss = self.lexicon_loss(bivalency_array,
                                                   attention_weights, text_mask)
                loss = loss + (self.bivalency_regularizer * bivalency_loss)

            for metric in self.metrics.values():
                metric(logits, label)

            output_dict["loss"] = loss

        return output_dict

    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
        return {metric_name: metric.get_metric(reset)
                for metric_name, metric in self.metrics.items()}

    def decode(self, output_dict: Dict[str, torch.Tensor]
               ) -> Dict[str, torch.Tensor]:
        predictions = output_dict['class_probabilities'].cpu().data.numpy()
        argmax_indices = np.argmax(predictions, axis=-1)
        labels = [self.vocab.get_token_from_index(x, namespace="labels")
                  for x in argmax_indices]
        output_dict['label'] = labels
        return output_dict


# http://codecombat.com/play/level/diamond-dozen

# ADVICE - use fast sword

# Claim the coins while defeating the marauding ogres.
# If you defeat the ogre with the most health, the rest of the ogres will run!
# Coins vanish quickly after appearing, so be sure to find the best value!

def findMostHealth(enemies):
    target = None
    targetHealth = 0
    enemyIndex = 0
    while enemyIndex < len(enemies):
        enemy = enemies[enemyIndex]
        if enemy.health > targetHealth:
            target = enemy
            targetHealth = enemy.health
        enemyIndex += 1
    return target

# Make a function named findOptimalCoin which returns the coin with the best value.
# Coins rapidly appear and disappear, so pick the best coin.
# Optimize your path by going for the coin with the largest value over distance.
def findOptimalCoin(coins):
    optimal = -1000
    thione = None
    for coin in coins:
        if coin.value / hero.distanceTo(coin) > optimal:
            optimal = coin.value / hero.distanceTo(coin)
            thione = coin
    return thione

while True:
    enemies = hero.findEnemies()
    enemy = findMostHealth(enemies)
    if enemy and enemy.health > 15:
        while enemy.health > 0:
            hero.attack(enemy)
    else:
        coins = hero.findItems()
        coin = None
        coin = findOptimalCoin(coins)  # ∆ Uncomment this once you've written the function.
        if coin:
            if hero.isReady('jump'):  # jump if you hero to slow
                hero.jumpTo(coin)  # also your hero can be too fast and grub too much coins
            else:  # use week but fast sword. Use best boots with jumps or ring of speed, but not both of them
                hero.moveXY(coin.pos.x, coin.pos.y)


#Locating Restriction Sites
#position and length of reverse palindrome

fasta = """>Rosalind_9452
ATATCAACTGTGGGCAAGATACTTCAAGCGTTTGATACAGTTGCCTCCGCCTTAGGCGCA
TACATCCAACAAAGGCCGCGCGTGGCCAATGAGGAATGAGGAGTTCCAGGCTAGAACAAA
CTTGACTAGGTGCTGACTATAGGCCTTGCCCGACACGGAACAGGGACCTCATTTACTGGA
TGCCTCCGATGCTCCGAATCGTGAAAGTCTGATAGACCAACATATTAATCGAACGAGGGG
ACAGGTCTTCTCAAACTCTCTGCCTTAGAGGCCCCTTATAGCTAGGTATCTGGGATGAAG
CGTATTACGAACCCGCAAATGTCATGATACTGAACCGTGGGACTGCAGATCCTGAGTGCC
CTCCCAGAACCCATCCTTTACATACTCAGCCTTTATAGTTAGAGTTCTCGGGTCGCTGGC
GTGTGGCCAAACCTAGAAGAAATCCAAAAGTGTCCGAATTCGGACCTGCTAGAACGCCAG
TAACGGATGGCGAAACACCAGCCACTGTACATCGGGGGCATTCCCTTTCTTGTGCCCTTC
CGTGAAAACGTTAGAACGGTTCTCCTAATATGGATTCGTTGAGGCTTCCGGCTCCTGAAT
TAGGGTTATTGAAAATATTGTGTATCGACACATTATAACATTAAAAGGTGCTAGAAAACG
ATATGAGTTTCAATCCCGGCGCGCTATGACTCTGTCGGACTCCTCACGGTTAGTTGCTTT
CATGTTAAATATTCAGCAGATAAATACGGTATCGGCAAGCCAGCTTATGCACGGCTCTTC
CCGACTTGAACACAGGAGAACCGCAACCAGGGTCTTAGCCGAGGACCATATTCTTCATTG
CCACAACGGGGCCTAACCGCCACAAGTCTCAATCAGACATTAGACCCGGGATTGGGATCG
GCCGATTCCGGGGCGC"""

def reverse_complement(dna):
    DNA = dna.upper()
    DNA_list = list(DNA)
    complement_list = []
    for i in DNA_list:
        if i == 'A':
            complement_list.append('T')
        elif i == 'T':
            complement_list.append('A')
        elif i == 'G':
            complement_list.append('C')
        elif i == 'C':
            complement_list.append('G')
    reverse_complement_list = complement_list
    reverse_complement_list.reverse()
    s = ''
    reverse_complement = s.join(reverse_complement_list)
    return reverse_complement

def fasta_into_dict(fasta):
    fasta_list = fasta.split()
    fasta_dict = {}
    for line in fasta_list:
        if line[0] == '>':
            header = line[1:]
            fasta_dict[header] = []
        else:
            fasta_dict[header].append(line[:])  #creates list within value {'header': ['ABS', 'ABC']}
    for header in fasta_dict:
        fasta_dict[header] = ''.join(fasta_dict[header])
    return fasta_dict

fasta_dict = fasta_into_dict(fasta)

for header, sequence in fasta_dict.items():         #loop over sequences in fasta file (here not necessary)
    for i in range(len(sequence)):                  #iterate over every base in sequence
        for a in range(2, 7):                       #reverse palindrome of length 4-12
            first = sequence[i : i + a]             #first half of palindrome
            second = sequence[i + a : i + a * 2]    #second half
            reversec = reverse_complement(second)   #reverse complement of second half
            if first == reversec:                   #check if identical
                print(i+1, a*2)                     #print position and length


#Ingresar nombres en una lista, luego buscar un nombre y de encontrarlo decir en qué posición está.

cant = int(input("¿Cuántos nombres quiere ingresar? -> "))
lista_nombres = ["" for x in range(cant)]

for i in range(len(lista_nombres)):
    nombre = input("Ingrese un nombre: ")
    lista_nombres[i] = nombre

print(lista_nombres)

#====|| Buscar nombre y posición ||=======
nombreBuscado = input("Ingrese un nombre a buscar: ")
for i in range(len(lista_nombres)):
    if nombreBuscado == lista_nombres[i]:
        posicion = i
        print(f"El nombre {nombreBuscado} se encuentra en la posición {posicion}")
if nombreBuscado not in lista_nombres:
    print("El nombre buscado no se encuentra en la lista")


# Definition for singly-linked list.
# class ListNode(object):
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution(object):
    def addTwoNumbers(self, l1, l2):
        """
        :type l1: ListNode
        :type l2: ListNode
        :rtype: ListNode
        """
        n_1 = ''
        n_2 = ''

        while l1 is not None:
            n_1 += str(l1.val)
            l1 = l1.next
        while l2 is not None:
            n_2 += str(l2.val)
            l2 = l2.next

        res = str(int(n_1[::-1]) + int(n_2[::-1]))[::-1]
        head = ListNode(res[0])
        cur_node = head

        for r in res[1:]:
            cur_node.next = ListNode(r)
            cur_node = cur_node.next
        return head


from collections import defaultdict

skills = list(map(str,input().split()))
k = len(skills)

n = int(input())

# 후행 스킬
adj = [[] for _ in range(k)]

# 선행스킬의 갯수 -> 0 이면 단독스킬
ind = [0 for _ in range(k)]

# 스킬을 인덱스 번호로
def conv(skill_name):
    return skills.index(skill_name)

for _ in range(n):
    s,e = map(str,input().split())
    s,e = conv(s),conv(e)
    adj[s].append(e)
    ind[e] += 1

def dfs(idx,path):

    if len(adj[idx]) == 0:
        print(' '.join(path))
        return
    for nxt in adj[idx]:
        dfs(nxt,path+[skills[nxt]])

for i in range(k):
    if ind[i] == 0:
        dfs(i,[skills[i]])


# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    def removeElements(self, head: Optional[ListNode], val: int) -> Optional[ListNode]:
        dummy=ListNode(next=head)
        prev=dummy # Because the element we have to remove might be the first element
        curr=prev.next
        while curr:
            if curr.val==val:
                prev.next=curr.next
                curr=curr.next
            else:
                prev=curr
                curr=curr.next
        return dummy.next


from . import DataTypes
import random
import numpy as np
import pandas as pd
import re

class Codon:
    def __init__(self) -> None:
        self.__codonTable = {
            'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
            'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
            'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
            'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
            'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
            'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
            'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
            'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
            'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
            'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
            'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
            'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
            'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
            'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
            'TAC':'Y', 'TAT':'Y', 'TAA':'STOP', 'TAG':'STOP',
            'TGC':'C', 'TGT':'C', 'TGA':'STOP', 'TGG':'W',
        }

    def DNAtoAA(self, DNA):
        return self.__codonTable[DNA]

    def AAtoDNA(self,AA):
        key_list = list(self.__codonTable.keys())
        val_list = list(self.__codonTable.values())
        position = val_list.index(AA)
        return(key_list[position])

    def GetRandomAA(self):
        val_list = list(self.__codonTable.values())
        return(random.sample(val_list,1)[0])

def GenerateProteinSyntheticFasta(variants):
    #with open ("data.tsv", 'w') as f:
    #    for v in proteinVariants:
    #        f.write(v.Export())
    #    for v in rnaVariants:
    #        f.write(v.Export())

    #what do we want?
    #lets say we spike in total 99 AMR genes. 33wt, 33res, 33other
    #randomly choose 99 unique AROs.
    #for wt: make sure the base is wt then insert it into fasta.
    #for mut: make sure the base is mut then insert it
    #   additional: ensure all mutation categories (i.e. fs, single, multi) are in there.
    #for other: make sure the base is not wt or mut then insert it.

    codonTable = Codon()
    fasta = []
    uniqueARO = list(set([x.aro.GetARO() for x in variants]))
    if ("3003686" in uniqueARO):
        uniqueARO.remove("3003686")
    random.seed(25041)

    variantMutationTypes = list(set([x.GetMutType() for x in variants]))

    variantSets = {}
    n = 40
    random.shuffle(uniqueARO)
    missing = 0

    for type in variantMutationTypes:
        variantSets[type] = []
        for i in range(missing + int(n/len(variantMutationTypes))):
            added = False
            for aro in uniqueARO:
                seqList = [x for x in variants if x.aro.GetARO() == aro]
                #print(seq.GetMutType() + " : " + type)
                for seq in seqList:
                    if (seq.GetMutType() == type):
                        variantSets[type].append(aro)
                        uniqueARO.remove(aro)
                        added = True
                        break
                if (added):
                    break
        if (len(variantSets[type]) != int(n/len(variantMutationTypes))):
            for i in range(int(n/len(variantMutationTypes)) - len(variantSets[type])):
                added = False
                for aro in uniqueARO:
                    seqList = [x for x in variants if x.aro.GetARO() == aro]
                    #print(seq.GetMutType() + " : " + type)
                    for seq in seqList:
                        if (seq.GetMutType() == "Single"):
                            variantSets[type].append(aro)
                            uniqueARO.remove(aro)
                            added = True
                            break
                    if (added):
                        break
    variantSets["wt"] = uniqueARO[-n:]
    del uniqueARO[-n:]
    variantSets["other"] = uniqueARO[-n:]
    del uniqueARO[-n:]

    for type in variantSets.keys():
        for aro in variantSets[type]:
            variantObjects = [x for x in variants if x.aro.GetARO() == aro]
            if (type == "wt"):
                #take the first object with given ARO
                seq = variantObjects[0]
            elif (type == "other"):
                #again, take the first object with given ARO
                seq = variantObjects[0]
            else:
                #mutants, scan through objects until the muttype matches, return that match
                for vo in variantObjects:
                    if (vo.GetMutType() == type):
                        seq = vo
                        break

            dna = seq.aro.dna
            mutType = seq.mutType
            if (type == "wt"):
                label = "wt"
            elif (type == "other"):
                label = "other"
            else:
                label = "mut"
            header = ">" + str(seq.aro.aro) + "|protein|" + label + "|" + str(mutType) + "|"

            for snp in seq.snp:
                if (type == 'wt'):
                    replacement = codonTable.AAtoDNA(snp.wt)
                elif (type == 'other'):
                    aa = codonTable.GetRandomAA()
                    while (aa == snp.mut or aa == snp.wt):
                        aa = codonTable.GetRandomAA()
                    replacement = codonTable.AAtoDNA(aa)
                else: # (type == 'mut'):
                    if (snp.mut == "fs"):
                        replacement = codonTable.AAtoDNA(snp.wt)
                    elif (snp.mut == "*"):
                        replacement = codonTable.AAtoDNA("STOP")
                    else:
                        replacement = codonTable.AAtoDNA(snp.mut)

                dna = mutType.GenerateSyntheticDNAForProtein(dna, replacement, snp.position)
                header = header + str(snp.position) + ":" + codonTable.DNAtoAA(replacement) + ";"
            fasta.append(header + "\n" + dna)
    with open ("SyntheticProteinVariants.fasta", "w") as f:
        f.write("\n".join(fasta))

def GenerateRRNASyntheticFasta(variants):

    #what do we want?
    #lets say we spike in total 99 AMR genes. 33wt, 33res, 33other
    #randomly choose 99 unique AROs.
    #for wt: make sure the base is wt then insert it into fasta.
    #for mut: make sure the base is mut then insert it
    #   additional: ensure all mutation categories (i.e. fs, single, multi) are in there.
    #for other: make sure the base is not wt or mut then insert it.

    codonTable = Codon()
    fasta = []
    uniqueARO = list(set([x.aro.GetARO() for x in variants]))
    if ("3003686" in uniqueARO):
        uniqueARO.remove("3003686")
    random.seed(25041)

    variantMutationTypes = list(set([x.GetMutType() for x in variants]))

    variantSets = {}
    n = 40
    random.shuffle(uniqueARO)
    missing = 0

    for type in variantMutationTypes:
        variantSets[type] = []
        for i in range(missing + int(n/len(variantMutationTypes))):
            added = False
            for aro in uniqueARO:
                seqList = [x for x in variants if x.aro.GetARO() == aro]
                #print(seq.GetMutType() + " : " + type)
                for seq in seqList:
                    if (seq.GetMutType() == type):
                        variantSets[type].append(aro)
                        uniqueARO.remove(aro)
                        added = True
                        break
                if (added):
                    break
        if (len(variantSets[type]) != int(n/len(variantMutationTypes))):
            for i in range(int(n/len(variantMutationTypes)) - len(variantSets[type])):
                added = False
                for aro in uniqueARO:
                    seqList = [x for x in variants if x.aro.GetARO() == aro]
                    #print(seq.GetMutType() + " : " + type)
                    for seq in seqList:
                        if (seq.GetMutType() == "Single"):
                            variantSets[type].append(aro)
                            uniqueARO.remove(aro)
                            added = True
                            break
                    if (added):
                        break
    variantSets["wt"] = uniqueARO[-n:]
    del uniqueARO[-n:]
    variantSets["other"] = uniqueARO[-n:]
    del uniqueARO[-n:]

    for type in variantSets.keys():
        for aro in variantSets[type]:
            variantObjects = [x for x in variants if x.aro.GetARO() == aro]
            if (type == "wt"):
                #take the first object with given ARO
                seq = variantObjects[0]
            elif (type == "other"):
                #again, take the first object with given ARO
                seq = variantObjects[0]
            else:
                #mutants, scan through objects until the muttype matches, return that match
                for vo in variantObjects:
                    if (vo.GetMutType() == type):
                        seq = vo
                        break

            dna = seq.aro.dna
            mutType = seq.mutType
            if (type == "wt"):
                label = "wt"
            elif (type == "other"):
                label = "other"
            else:
                label = "mut"
            header = ">" + str(seq.aro.aro) + "|rna|" + label + "|" + str(mutType) + "|"

            for snp in seq.snp:
                if (type == 'wt'):
                    replacement = snp.wt
                elif (type == 'other'):
                    replacement = "A"
                    while (replacement == snp.mut or replacement == snp.wt):
                        replacement = random.sample(["A","T","G","C"],1)[0]
                else:
                    replacement = snp.mut
                header = header + str(snp.position) + ":" +replacement + ";"
                dna = mutType.GenerateSyntheticDNAForRNA(dna, replacement, snp.position)

                #codon = codonTable.AAtoDNA(aa)
                #start = (snp.position -1) *3
                #end = (snp.position -1) *3 + 3
                #originalDNA = dna[start:end]
                #originalAA = codonTable.DNAtoAA(originalDNA)
                #dna = dna[:start] + codon + dna[end:]
            fasta.append(header + "\n" + dna)

    with open ("SyntheticRRnaVariants.fasta", "w") as f:
        f.write("\n".join(fasta))

def CheckAccuracy(resultTable, fastaPath):
    variants = resultTable

    with open(fastaPath, 'r') as f:
        fasta = f.readlines()
    fasta = [x for x in fasta if (x.startswith(">") and x.find("other") == -1)]
    #variants = [x for x in variants if (x.find("Resistant") > -1)]

    referenceVariants = {}
    for line in fasta:
        l = line.strip().replace(">","").split("|")
        aro = l[0]
        mutationType = l[2]
        mutationClass = l[3]
        snp = l[4]
        if aro in referenceVariants.keys():
            referenceVariants[aro].append([mutationType, mutationClass, snp])
        else:
            referenceVariants[aro] = []
            referenceVariants[aro].append([mutationType, mutationClass, snp])

    detectedVariants = {}
    for line in variants:
        l = line.strip().split("\t")
        aro = l[0]
        mutationType = l[2]
        mutationClassification = l[3]
        snp = l[4]
        if aro in detectedVariants.keys():
            detectedVariants[aro].append([mutationType, mutationClassification, snp])
        else:
            detectedVariants[aro] = []
            detectedVariants[aro].append([mutationType, mutationClassification, snp])

    keyDictionary = {"wt": 0, "mut": 1, "other":0,
                    "Wildtype" : 0, "Resistant Variant" : 1, "Other Variant":0, "Not Found":0, "Partial":0, "True":1, "False":0}

    matrix = np.zeros ((3, 5))
    matrixAro = np.empty((3,5), dtype=object)

    matrixBinary = np.zeros((2,2))
    matrixBinaryAro =  np.empty((2,2), dtype=object)
    matrixBinaryAro[0,0] = ""
    matrixBinaryAro[0,1] = ""
    matrixBinaryAro[1,0] = ""
    matrixBinaryAro[1,1] = ""
    sum = 0
    for refARO in referenceVariants.keys():
        #refARO is a list
        referenceVariants[refARO] = [list(x) for x in set(tuple(x) for x in referenceVariants[refARO] )]
        for refVariant in referenceVariants[refARO]:
            refPosition = list(set(re.findall( '(\d+)', refVariant[2])))
            if ("531" in refPosition):
                print("test")
            if (refARO in detectedVariants.keys()):
                detected = [x for x in detectedVariants[refARO] if list(set(re.findall( '(\d+)', x[2]))) == refPosition]
            else:
                detected = []

            if (len(detected) > 1):
                expected = [x.split(":")[1] for x in refVariant[2].strip().split(";") if x]
                expected = [x.replace("U", "T") for x in expected]
                expected.sort()
                found = False
                for i in detected:
                    bases = [x[-1].upper() for x in i[2].split(";")]
                    bases = [x.replace("*", "STOP") for x in bases]
                    bases.sort()
                    if expected == bases:
                        matrixBinary[keyDictionary[refVariant[0]],keyDictionary[i[1]]] += 1
                        label = refARO + "|" + refVariant[2] + "|" + detectedVariant[2] #+ "|" + detectedVariant[3]
                        matrixBinaryAro[keyDictionary[refVariant[0]],keyDictionary[i[1]]] = matrixBinaryAro[keyDictionary[refVariant[0]],keyDictionary[i[1]]] + str(label) + ";"
                        found = True
                if(found == False):
                    #likely a wildtype because the mutant base is not found.
                    if (refVariant[0]=="wt"):
                        matrixBinary[keyDictionary[refVariant[0]],0] += 1
                        label = refARO + "|" + refVariant[2] + "|" + detectedVariant[2] #+ "|" + detectedVariant[3]
                        matrixBinaryAro[keyDictionary[refVariant[0]],0] = matrixBinaryAro[keyDictionary[refVariant[0]],0] + str(label) + ";"
                    else:
                        #a resist type not ID'd
                        matrixBinary[keyDictionary[refVariant[0]],0] += 1
                        label = refARO + "|" + refVariant[2] + "|" + detectedVariant[2] #+ "|" + detectedVariant[3]
                        matrixBinaryAro[keyDictionary[refVariant[0]],0] = matrixBinaryAro[keyDictionary[refVariant[0]],0] + str(label) + ";"

                        #print(refVariant)
                        #print(detected)
            elif(len(detected) == 1):
                detectedVariant = detected[0]
                matrixBinary[keyDictionary[refVariant[0]],keyDictionary[detectedVariant[1]]] += 1
                label = refARO + "|" + refVariant[2] + "|" + detectedVariant[2] #+ "|" + detectedVariant[3]
                matrixBinaryAro[keyDictionary[refVariant[0]],keyDictionary[detectedVariant[1]]] = matrixBinaryAro[keyDictionary[refVariant[0]],keyDictionary[detectedVariant[1]]] + str(label) + ";"
            else:
                print("notfound")
            sum += 1
    totalRes = len([x for x in referenceVariants.values() if x[0][0] == "mut"])
    totalGroundTruth = sum

    #df = pd.DataFrame(matrixBinary)
    #df.columns = ["Non-Resistant", "Resistant"]
    #df["Truth\\Pred"] = ["Non-Resistant", "Resistant"]
    #df = df.set_index("Truth\\Pred")
    #df.to_csv("accuracy.tsv", sep ='\t', header = True , index = True)

    df = pd.DataFrame(matrixBinaryAro)
    df.columns = ["Non-Resistant", "Resistant"]
    df["Truth\\Pred"] = ["Non-Resistant", "Resistant"]
    df = df.set_index("Truth\\Pred")
    #df.to_csv("accuracy_aro.tsv", sep ='\t', header = True , index = True)

    matrix = AccuracyMetrics(matrixBinary[1,1], matrixBinary[0,0], matrixBinary[0,1], matrixBinary[1,0])

    return matrix, totalRes, totalGroundTruth

#cardJsonPath = "/mnt/d/OneDrive/ProjectAMRSAGE/AMR_Metagenome_Simulator-master/full/card/card.json"
#cardSnpPath = "/mnt/d/OneDrive/ProjectAMRSAGE/AMR_Metagenome_Simulator-master/full/card/snps.txt"

#CARD = ParseCardData.ParseJson(cardJsonPath)
#proteinVariants, rnaVariants = ParseCardData.ParseSNP(cardSnpPath, CARD)

#GenerateProteinSyntheticFasta(proteinVariants)
#GenerateRRNASyntheticFasta(rnaVariants)

#resultPath = "./syntheticBalanced_DetectedVariants.tsv"
#with open (resultPath , 'r') as f:
#    variants = f.readlines()[1:]
#fastaPath = "./sim/SyntheticProteinVariants.fasta"
#test = CheckAccuracy(variants, fastaPath)
#print(test)
#50 protein sequences of each of the following:
# 1) wt
# 2) mut
# 3) other
#matrix, totalRes, total = CheckAccuracy("synthetic.tsv", "SyntheticVariants.fasta")


"""
50. Pow(x, n)

# x^n --> x2 ^ n//2 ....

logn
"""
def myPow(self, x: float, n: int) -> float:
    # n negative case
    if n < 0:
        x = 1 / x
        n = -n
    # divide by n//2 make x^2
    ans = 1
    while n:
        if n % 2:
            ans *= x
        x *= x
        n = n // 2
    return ans


#!/usr/bin/python3
"""
This module takes in a URL,
sends a request to the URL and
displays the value of the 'X-Request-Id' variable
found in the header of the response.
"""

if __name__ == '__main__':
    from urllib.request import urlopen
    import sys
    url = sys.argv[1]
    with urlopen(url) as response:
        response_headers = response.info()
        print(response_headers.get('X-Request-Id'))


import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from generator import *
import math

# A simple fork-join workflow

def forkjoin(N, alpha, runtimeDist, sizeDist):
    # Width of workflow
    W = int(math.ceil(N/float(alpha)))

    # Number of levels
    L = int(math.ceil(N/float(W+1)))

    w = Workflow(name="forkjoin", description="""Fork-Join workflow (Figure 2b in Rahman et al)""")

    lastlevel = [File(name="infile.dat", size=sizeDist()*GB)]

    for i in range(0,L):
        fork = Job(id="task_%d"%i, namespace="parallel", name="Task", runtime=runtimeDist()*SECONDS, inputs=lastlevel)
        w.addJob(fork)

        lastlevel = []
        for j in range(0,W):
            tin = File(name="task_%d_%d_in.dat"%(i,j), size=sizeDist()*GB)
            tout = File(name="task_%d_%d_out.dat"%(i,j), size=sizeDist()*GB)

            fork.addOutput(tin)
            lastlevel.append(tout)

            j = Job(id="task_%d_%d"%(i,j), namespace="parallel", name="Task", runtime=runtimeDist()*SECONDS, inputs=[tin], outputs=[tout])
            w.addJob(j)

    outputs = [File(name="outfile.dat", size=sizeDist()*GB)]
    join = Job(id="join", namespace="parallel", name="Task", runtime=runtimeDist()*SECONDS, inputs=lastlevel, outputs=outputs)
    w.addJob(join)

    return w

def main(*args):
    class ForkJoin(Main):
        def setoptions(self, parser):
            self.parser.add_option("-N", "--numjobs", dest="N", metavar="N", type="int", default=50,
                help="Number of jobs in workflow [default: %default]")
            self.parser.add_option("-a", "--alpha", dest="alpha", metavar="A", type="float", default=5,
                help="Ratio of number of jobs to width (e.g. W = ceil(N/alpha)) [default: %default]")

            self.parser.add_option("", "--rtlow", dest="rtlow", metavar="t", type="float", default=100,
                help="Lower bound on runtime in seconds [default: %default]")
            self.parser.add_option("", "--rthigh", dest="rthigh", metavar="t", type="float", default=500,
                help="Upper bound on runtime in seconds [default: %default]")

            self.parser.add_option("", "--slow", dest="slow", metavar="t", type="float", default=1,
                help="Lower bound on file size in GB [default: %default]")
            self.parser.add_option("", "--shigh", dest="shigh", metavar="t", type="float", default=5,
                help="Upper bound on file size in GB [default: %default]")

        def genworkflow(self, options):
            return forkjoin(options.N, options.alpha,
                UniformDistribution(options.rtlow, options.rthigh),
                UniformDistribution(options.slow, options.shigh))

    ForkJoin().main(*args)

if __name__ == '__main__':
    main(*sys.argv[1:])


#!/usr/bin/python3
"""
Module 100-append_after
Define append_after
searchs for a string and adds a text
"""

def append_after(filename="", search_string="", new_string=""):
    """Inserts a line in a text file"""
    content = []
    with open(filename, mode="r") as f:
        for line in f:
            content.append(line)
        f.seek(0)

    for i in range(len(content)):
        if search_string in content[i]:
            if i == len(content) - 1:
                content.append(new_string)
            else:
                content.insert(i + 1, new_string)

    with open(filename, mode="w") as f:
        for line in content:
            f.write(line)


import os
import re
import mmap
from pymongo import MongoClient

import config

client = MongoClient()
db = client[config.MONGO_DATABASE]
sources_collection = db[config.MONGO_COLLECTION_SOURCES]
countries_collection = db[config.MONGO_COLLECTION_COUNTRIES]

minerals_regex = re.compile(bytes("(%s)" % "|".join(config.MINERALS), 'utf-8'), re.IGNORECASE | re.MULTILINE)

for source in sources_collection.find({'newest_xls_local_csv_directory': {'$exists': 1}}):
    print("Processing source %s..." % source['_id'])
    for table in os.listdir(source['newest_xls_local_csv_directory']):
        print("  └ %s..." % table)
        with open(os.path.join(source['newest_xls_local_csv_directory'], table), 'rb', 0) as file, \
                mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as file_mmap:
            mineral_matches = [mineral.decode('utf-8').lower() for mineral in minerals_regex.findall(file_mmap)]
            mineral_matches = list(dict.fromkeys(mineral_matches))  # deduplicate
            print('   ', mineral_matches)
            sources_collection.update_one(
                {'_id': source['_id']},
                {'$set': {'minerals': mineral_matches}}
            )
            countries_collection.update_many(
                {'_id': {'$in': source['used_by']}},
                {'$addToSet': {'minerals': {'$each': mineral_matches}}}
            )


from sklearn.metrics import f1_score
from simpletransformers.classification import MultiLabelClassificationModel

import sys
import os

import pandas as pd
import datetime
import csv
import numpy as np

root_path = os.path.dirname(os.path.dirname(
    os.path.dirname(os.path.abspath(__file__))))

prep_path = os.path.join(root_path, 'preprocess')
sys.path.append(prep_path)

split_path = os.path.join(root_path, 'dataset/split')

curr_path = os.path.dirname(os.path.abspath(__file__))

from reader import Reader

def f1_evaluate(true, pred):
    for p in pred:
        for i in range(len(p)):
            if p[i] >= 0.5:
                p[i] = 1
            else:
                p[i] = 0

    score = f1_score(true, pred, average = 'macro')
    label = f1_score(true, pred, average = None)
    print(score)
    print(label)

    return score

class BERTMultilabel:
    def __init__(self):
        self.dummy = None
        self.reader = Reader()

    def train_model(self, training_data, training_args, base_model_path, test_data):
        model = MultiLabelClassificationModel(
            "bert",
            base_model_path,
            num_labels=8,
            use_cuda = False, # Highly recommended to set use_cuda = True to ultilize GPU (if available) for training
            args = training_args,
        )

        temp = [self.reader.assign_label(s) for s in training_data]

        test_data = [self.reader.assign_label(s) for s in test_data]
        test_text = [s[0] for s in test_data]
        test_label = [s[1] for s in test_data]

        eval_df = pd.DataFrame(test_data, columns=['text', 'labels'])

        model.train_model(pd.DataFrame(temp, columns=['text', 'labels']), eval_df = eval_df, f1_macro = f1_evaluate)

        return model

    def test_model(self, model, test_data):
        if model == None:
            print('No model found')
            return

        test_data = [self.reader.assign_label(s) for s in test_data]
        test_text = [s[0] for s in test_data]
        test_label = [s[1] for s in test_data]

        predictions, dummy = model.predict(test_text)
        pred = pd.DataFrame(predictions)
        pd_label = pd.DataFrame(test_label)

        np.set_printoptions(formatter={'float': '{: 0.3f}'.format})

        f1_macro = f1_score(pd_label, pred, average='macro')
        f1_all = f1_score(pd_label, pred, average=None)

        print(f1_macro)
        print(f1_all)

        return f1_macro, f1_all

# # Training args
# args = {"reprocess_input_data": True,
# "overwrite_output_dir": True,
# "num_train_epochs": 30,
# 'fp16': False,
# "use_early_stopping": True,
# 'learning_rate': 4e-5,
# 'evaluate_during_training' : True,
# 'early_stopping_metric' : 'f1_macro',
# 'early_stopping_metric_minimize': False,
# 'save_model_every_epoch' : False,
# 'train_batch_size' : 8
# }

# # Read the splits

# r = Reader()

# split_0 = r.read_from_split(split_path + '/split_0.tsv')
# split_1 = r.read_from_split(split_path + '/split_1.tsv')

# base_model_path = 'bert-large-uncased'

# bm = BERTMultilabel()

# # Train 01 base model
# train = split_0
# test = split_1

# model = bm.train_model(train, args, base_model_path, test)


# plugin_pose_DeepLabCut.py

from qtpy.QtCore import Qt, QPointF
from qtpy.QtGui import QColor, QPainter, QPen, QPolygonF
from qtpy.QtWidgets import QMessageBox, QWidget
from pose.pose import PoseBase
from utils import get_colormap
import csv
import numpy as np
import h5py
from os.path import splitext

"""
Bento plugin that provides support for MARS-style pose files
represented MatLab .mat files.

Implements a class derived from PoseBase, which is an abstract class.
"""

class PoseDLCBase(PoseBase):
    """
    Base class for DeepLabCut pose support

    Includes implementation common to generic and mouse-specific DLC files
    """

    def __init__(self):
        super().__init__()

    def getFileSearchPattern(self) -> str:
        return "*.h5 *.csv"

    def _validateFileH5(self, parent_widget, file_path: str) -> bool:
        h5 = h5py.File(file_path, 'r')
        default_group = h5[list(h5.keys())[0]]
        if 'table' not in default_group.keys():
            QMessageBox.warning(parent_widget, "Add Pose ...", "No pose data found in pose file")
            return False
        return True

    def _validateFileCSV(self, parent_widget, file_path: str) -> bool:
        result = True
        with open(file_path, 'r') as csvFile:
            reader = csv.reader(csvFile)
            header1 = next(reader)
            if (not isinstance(header1[1], str) or
                header1[1] != 'bodyparts' or
                (len(header1) - 2) % 3 != 0):   # "", "bodyparts", 3 x <bodypart_name>, ...
                result = False
            csvFile.seek(0)
        return result

    def validateFile(self, parent_widget: QWidget, file_path: str) -> bool:
        """
        Default implementation does no checking,
        but we can do better than that
        """
        _, ext = splitext(file_path)
        ext = ext.lower()
        if ext == '.h5':
            return self._validateFileH5(parent_widget, file_path)
        elif ext == '.csv':
            return self._validateFileCSV(parent_widget, file_path)
        else:
            QMessageBox.warning(parent_widget, "Extension not supported",
                f"The file extension {ext} is not supported.")

    def _loadPoses_h5(self, parent_widget: QWidget, path: str):
        raise NotImplementedError("Please implement this in your derived class")

    def _loadPoses_csv(self, parent_widget: QWidget, path: str):
        raise NotImplementedError("Please implement this in your derived class")

    def loadPoses(self, parent_widget: QWidget, path: str):
        _, self.file_extension = splitext(path)
        self.file_extension = self.file_extension.lower()
        if self.file_extension == '.h5':
            self._loadPoses_h5(parent_widget, path)
        elif self.file_extension == '.csv':
            self._loadPoses_csv(parent_widget, path)
        else:
            QMessageBox.warning(parent_widget, "Extension not supported",
                f"The file extension {self.file_extension} is not supported.")

class PoseDLC_generic(PoseDLCBase):

    def __init__(self):
        super().__init__()
        self.file_extension = None
        self.frame_points = []
        self.num_frames = 0
        # Construct colors based on the number of body parts and the color map
        self.colormap_data = get_colormap('turbo')
        self.pose_colors = None

    def generatePoseColors(self, n):
        if n < 1 or n > 256:
            raise ValueError("generatePoseColors: n must be between 1 and 256 inclusive")
        if n == 1:
            self.pose_colors = [QColor(self.colormap_data[127])]
        else:
            step = 255. / (n - 1)
            self.pose_colors = [QColor(self.colormap_data[round(ix * step)]) for ix in range(n)]

    def drawPoses(self, painter: QPainter, frame_ix: int):
        frame_ix = min(frame_ix, self.num_frames)
        painter.setPen(Qt.NoPen)
        if len(self.frame_points[frame_ix]) > len(self.pose_colors):
            raise ValueError("drawPoses: number of points t draw exceeds number of body parts")
        for ix, point in enumerate(self.frame_points[frame_ix]):
            painter.setBrush(self.pose_colors[ix])
            painter.drawEllipse(point, 5.0, 5.0)

    def getFileSearchDescription(self) -> str:
        return "DeepLabCut generic pose files"

    def getFileFormat(self) -> str:
        return "DeepLabCut_generic"

    def _loadPoses_h5(self, parent_widget, path: str):
        # TODO: need to change this from the mouse form to a more generic
        # version that just stores the points.
        h5 = h5py.File(path, 'r')
        default_group = h5[list(h5.keys())[0]]
        table = default_group['table']
        self.num_frames = len(table)
        isColorsGenerated = False
        for frame_ix in range(self.num_frames):
            # DLC stores each frame as (frame_idx, (pose_array))
            # we don't care about the frame_idx
            this_frame_data = table[frame_ix][1]
            this_frame_points = []
            vals_per_pt = 3
            num_points = int(len(this_frame_data) / vals_per_pt)
            if not isColorsGenerated:
                self.generatePoseColors(num_points)
            for point_ix in range(num_points):
                # each pose point has x, y and confidence
                # we don't care about confidence
                pt_x_ix = (vals_per_pt * point_ix) + 0
                pt_y_ix = (vals_per_pt * point_ix) + 1
                this_frame_points.append(QPointF(this_frame_data[pt_x_ix], this_frame_data[pt_y_ix]))
            self.frame_points.append(this_frame_points)

    def _loadPoses_csv(self, parent_widget, path: str):
        with open(path, 'r') as csvFile:
            vals_per_pt = 3 # x, y, confidence.  We don't care about confidence.
            reader = csv.reader(csvFile)
            header1 = next(reader)
            num_pts = int((len(header1) - 2) / vals_per_pt)
            self.generatePoseColors(num_pts)
            _ = next(reader)    # skip over second row of header
            while True:
                try:
                    row = next(reader)
                except StopIteration:
                    break
                this_frame_points = []
                for point_ix in range(num_pts):
                    pt_x_ix = (vals_per_pt * point_ix) + 2
                    pt_y_ix = (vals_per_pt * point_ix) + 3
                    this_frame_points.append(QPointF(float(row[pt_x_ix]), float(row[pt_y_ix])))
                self.frame_points.append(this_frame_points)
            self.num_frames = len(self.frame_points)

class PoseDLC_mouse(PoseDLCBase):

    def __init__(self):
        super().__init__()
        self.pose_polys = []
        self.num_frames = 0
        self.pose_colors = [Qt.blue, Qt.green]

    def drawPoses(self, painter: QPainter, frame_ix: int):
        frame_ix = min(frame_ix, self.num_frames)
        try:
            for mouse_ix in range(len(self.pose_polys[frame_ix])):
                painter.setPen(QPen(self.pose_colors[mouse_ix], 2.0))
                painter.drawPolyline(self.pose_polys[frame_ix][mouse_ix])
                painter.setBrush(Qt.red)
                painter.setPen(Qt.NoPen)
                painter.drawEllipse(self.pose_polys[frame_ix][mouse_ix].first(), 5.0, 5.0) # red dot on nose
        except IndexError as e:
            # number of video frames is greater than number of pose frames.  Oh well!
            pass

    def getFileSearchDescription(self) -> str:
        return "DeepLabCut mouse pose files"

    def getFileFormat(self) -> str:
        return "DeepLabCut_mouse"

    def _loadPoses_h5(self, parent_widget: QWidget, path: str):
        h5 = h5py.File(path, 'r')
        default_group = h5[list(h5.keys())[0]]
        table = default_group['table']
        self.pose_polys = []
        self.num_frames = len(table)
        for frame_ix in range(self.num_frames):
            # DLC stores each frame as (frame_idx, (pose_array))
            # we don't care about the frame_idx
            this_frame_data = table[frame_ix][1]
            frame_polys = []
            vals_per_pt = 3
            pts_per_mouse = 7
            vals_per_mouse = vals_per_pt * pts_per_mouse
            num_mice = int(len(this_frame_data) / vals_per_mouse)
            for mouse_ix in range(num_mice):
                # each pose point has x, y and confidence
                # we don't care about confidence
                poly = QPolygonF()
                pt_x_ix = (vals_per_mouse * mouse_ix) + 0
                pt_y_ix = (vals_per_mouse * mouse_ix) + 1
                nose = QPointF(this_frame_data[pt_x_ix], this_frame_data[pt_y_ix])
                poly.append(nose)
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (1 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (1 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (3 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (3 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (4 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (4 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (6 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (6 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (5 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (5 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (3 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (3 * vals_per_pt)]))
                poly.append(QPointF(
                    this_frame_data[pt_x_ix + (2 * vals_per_pt)],
                    this_frame_data[pt_y_ix + (2 * vals_per_pt)]))
                poly.append(nose)
                frame_polys.append(poly)
            self.pose_polys.append(frame_polys)

    def _loadPoses_csv(self, parent_widget: QWidget, path: str):
        with open(path, 'r') as csvFile:
            vals_per_pt = 3 # x, y, confidence.  We don't care about confidence.
            vals_per_mouse = 7
            reader = csv.reader(csvFile)
            header1 = next(reader)
            num_pts = int((len(header1) - 2) / vals_per_pt)
            if num_pts % vals_per_mouse != 0:
                QMessageBox.warning(parent_widget, "Improper number of points for mice",
                    f"Expected a multiple of 7 points, got {num_pts}")
                return
            _ = next(reader)    # skip over second row of header
            while True:
                try:
                    row = next(reader)
                except StopIteration:
                    break
                num_mice = int(num_pts / vals_per_mouse)
                frame_polys = []
                for mouse_ix in range(num_mice):
                    # each pose point has x, y and confidence
                    # we don't care about confidence
                    poly = QPolygonF()
                    pt_x_ix = (vals_per_mouse * mouse_ix) + 2
                    pt_y_ix = (vals_per_mouse * mouse_ix) + 3
                    nose = QPointF(row[pt_x_ix], row[pt_y_ix])
                    poly.append(nose)
                    poly.append(QPointF(
                        row[pt_x_ix + (1 * vals_per_pt)],
                        row[pt_y_ix + (1 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (3 * vals_per_pt)],
                        row[pt_y_ix + (3 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (4 * vals_per_pt)],
                        row[pt_y_ix + (4 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (6 * vals_per_pt)],
                        row[pt_y_ix + (6 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (5 * vals_per_pt)],
                        row[pt_y_ix + (5 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (3 * vals_per_pt)],
                        row[pt_y_ix + (3 * vals_per_pt)]))
                    poly.append(QPointF(
                        row[pt_x_ix + (2 * vals_per_pt)],
                        row[pt_y_ix + (2 * vals_per_pt)]))
                    poly.append(nose)
                    frame_polys.append(poly)
                self.pose_polys.append(frame_polys)
            self.num_frames = len(self.frame_points)

def register(registry):
    # construct and register the generic plugin
    pose_plugin = PoseDLC_generic()
    registry.register(pose_plugin.getFileFormat(), pose_plugin)

    # construct and register the MARS-style mouse-specific plugin
    pose_plugin = PoseDLC_mouse()
    registry.register(pose_plugin.getFileFormat(), pose_plugin)


1#!/usr/bin/env python
# -*- coding: utf-8 -*-
# =============================================================================
## @file ostap/tools/tests/test_tools_reweight3.py
#  Test for 3D-reweighting machinery
#  @author Vanya BELYAEV Ivan.Belyaev@itep.ru
#  @date 2023-01-20
# =============================================================================
"""Test for 3D-reweighting machinery
"""
# =============================================================================
__version__ = "$Revision:"
__author__  = "Vanya BELYAEV Ivan.Belyaev@itep.ru"
__date__    = "2023-01-20"
__all__     = ()  ## nothing to be imported
# =============================================================================
from   builtins               import range
from   ostap.core.pyrouts     import Ostap
import ostap.io.root_file
import ostap.parallel.kisa
import ostap.io.zipshelve     as     DBASE
import ostap.logger.table     as     T
from   ostap.histos.histos    import h1_axis, h2_axes, h3_axes
from   ostap.utils.timing     import timing
from   ostap.logger.colorized import attention, allright
from   ostap.plotting.canvas  import use_canvas
import ostap.logger.table     as     T
import ROOT, random, math, os, time
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__  or '__builtin__'  == __name__ :
    logger = getLogger ( 'ostap.test_tools_reweight3' )
else :
    logger = getLogger ( __name__ )
# =============================================================================
logger.info ( 'Test for 3D-Reweighting machinery')
# ============================================================================
from ostap.utils.cleanup import CleanUp
testdata   = CleanUp.tempfile ( suffix = '.root' , prefix ='ostap-test-tools-reweight3-' )

tag_data_3d = 'DATA_3D_histogram'

tag_data_xy = 'DATA_XY_histogram'
tag_data_yz = 'DATA_YZ_histogram'
tag_data_xz = 'DATA_XZ_histogram'

tag_data_x  = 'DATA_X_histogram'
tag_data_y  = 'DATA_Y_histogram'
tag_data_z  = 'DATA_Z_histogram'

tag_data    = 'DATA_tree'
tag_mc      = 'MC_tree'

dbname      = CleanUp.tempfile ( suffix = '.db' , prefix ='ostap-test-tools-reweight3-'   )

NDATA1      = 500000
NDATA2      = 400000
NDATA3      = 400000
NMC         = 500000

xmax        = 15.0
ymax        = 12.0
zmax        = 10.0

def prepare_data ( ) :
    #

    seed =  1234567890
    random.seed ( seed )
    logger.info ( 'Test *RANDOM* data will be generated/seed=%s' % seed  )
    ## prepare "data" histograms:

    # 1) 3D histogram
    ix , iy , iz  = 5 , 5 , 5
    h3_histo  = h3_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                          [ ymax/iy*i for i in range ( iy + 1 ) ] ,
                          [ zmax/iz*i for i in range ( iz + 1 ) ] )

    # 2) XY,YX &XZ histograms

    ix , iy , iz  =  7 ,  7 ,  7
    hxy_histo = h2_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                          [ ymax/iy*i for i in range ( iy + 1 ) ] )

    ix , iy , iz  =  8 ,  8 ,  8
    hyz_histo = h2_axes ( [ ymax/iy*i for i in range ( iy + 1 ) ] ,
                          [ zmax/iz*i for i in range ( iz + 1 ) ] )

    ix , iy , iz  =  9 ,  9 ,  9
    hxz_histo = h2_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                          [ zmax/iz*i for i in range ( iz + 1 ) ] )

    ## 3) X,Y,Z histogramss

    ix , iy , iz  = 100 , 100 , 100
    hx_histo  = h1_axis ( [ xmax/ix*i for i in range ( ix + 1 ) ] )
    hy_histo  = h1_axis ( [ ymax/iy*i for i in range ( iy + 1 ) ] )
    hz_histo  = h1_axis ( [ zmax/iz*i for i in range ( iz + 1 ) ] )

    with ROOT.TFile.Open ( testdata ,'recreate') as mc_file:

        mc_file.cd()

        datatree  = ROOT.TTree ( tag_data , 'data-tree' )
        datatree .SetDirectory ( mc_file )

        from array import array
        xvar = array    ( 'f', [0])
        yvar = array    ( 'f', [0])
        zvar = array    ( 'f', [0])
        datatree.Branch ( 'x' , xvar , 'x/F' )
        datatree.Branch ( 'y' , yvar , 'y/F' )
        datatree.Branch ( 'z' , zvar , 'z/F' )

        ## Gaussian 3D-component with correlations

        g3d = Ostap.Math.Gauss3D ( 0.1 * xmax  , 0.5 * ymax  , 0.8 * zmax  ,
                                   8           , 5           , 4           ,
                                   math.pi / 5 , math.pi / 5 , math.pi / 5 )

        for x, y , z in g3d.random ( NDATA1 , 0 , xmax , 0 , ymax , 0 , zmax ) :

            h3_histo .Fill ( x , y , z )

            hxy_histo.Fill ( x , y )
            hyz_histo.Fill ( y , z )
            hxz_histo.Fill ( x , z )

            hx_histo .Fill ( x )
            hy_histo .Fill ( y )
            hz_histo .Fill ( z )

            xvar [ 0 ] = x
            yvar [ 0 ] = y
            zvar [ 0 ] = z

            datatree.Fill()

        for i in range ( NDATA2 ) :

            x , y, z  = -1, -1, -1

            while not 0 < x < xmax : x = xmax - random.expovariate ( 1.0 / xmax )
            while not 0 < y < ymax : y =        random.expovariate ( 1.0 / ymax )
            while not 0 < z < zmax : z =        random.expovariate ( 1.0 / zmax )

            h3_histo .Fill ( x , y , z )

            hxy_histo.Fill ( x , y )
            hyz_histo.Fill ( y , z )
            hxz_histo.Fill ( x , z )

            hx_histo .Fill ( x )
            hy_histo .Fill ( y )
            hz_histo .Fill ( z )

            xvar [ 0 ] = x
            yvar [ 0 ] = y
            zvar [ 0 ] = z

            datatree.Fill()

        for i in range ( NDATA3 ) :

            x , y, z  = -1, -1, -1

            x = random.uniform ( 0 , xmax )
            y = random.uniform ( 0 , ymax )
            z = random.uniform ( 0 , zmax )

            h3_histo .Fill ( x , y , z )

            hxy_histo.Fill ( x , y )
            hyz_histo.Fill ( y , z )
            hxz_histo.Fill ( x , z )

            hx_histo .Fill ( x )
            hy_histo .Fill ( y )
            hz_histo .Fill ( z )

            xvar [ 0 ] = x
            yvar [ 0 ] = y
            zvar [ 0 ] = z

            datatree.Fill()

        datatree.Write()

        ## write the histogram
        mc_file [ tag_data_3d  ] = h3_histo

        mc_file [ tag_data_xy  ] = hxy_histo
        mc_file [ tag_data_yz  ] = hyz_histo
        mc_file [ tag_data_xz  ] = hxz_histo

        mc_file [ tag_data_x   ] = hx_histo
        mc_file [ tag_data_y   ] = hy_histo
        mc_file [ tag_data_z   ] = hz_histo

        ##
        mctree  = ROOT.TTree ( tag_mc , 'mc-tree' )
        mctree .SetDirectory ( mc_file )

        from array import array
        xvar = array  ( 'f', [ 0.0 ] )
        yvar = array  ( 'f', [ 0.0 ] )
        zvar = array  ( 'f', [ 0.0 ] )

        mctree.Branch ( 'x' , xvar , 'x/F' )
        mctree.Branch ( 'y' , yvar , 'y/F' )
        mctree.Branch ( 'z' , zvar , 'z/F' )

        for i in  range ( NMC ) :

            xv = random.uniform  ( 0 , xmax )
            yv = random.uniform  ( 0 , ymax )
            zv = random.uniform  ( 0 , zmax )

            xvar [ 0 ] = xv
            yvar [ 0 ] = yv
            zvar [ 0 ] = zv

            mctree.Fill()

        mctree.Write()

if not os.path.exists( testdata ) :
    with timing ( "Prepare input data" , logger = logger ) :
        prepare_data ()

# =============================================================================
## Read data from DB
# =============================================================================
with ROOT.TFile.open ( testdata , 'r' ) as dbroot :

    logger.info ( 'Test data is fetched from DBASE "%s"' % testdata )
    dbroot.ls_table ()

    h3d_data = dbroot [ tag_data_3d ] . clone()
    hxy_data = dbroot [ tag_data_xy ] . clone()
    hyz_data = dbroot [ tag_data_yz ] . clone()
    hxz_data = dbroot [ tag_data_xz ] . clone()
    hx_data  = dbroot [ tag_data_x  ] . clone()
    hy_data  = dbroot [ tag_data_y  ] . clone()
    hz_data  = dbroot [ tag_data_z  ] . clone()

# =============================================================================
## prebook MC histograms
# =============================================================================
ix , iy , iz  = 5 , 5 , 5
mc_3d = h3_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                  [ ymax/iy*i for i in range ( iy + 1 ) ] ,
                  [ zmax/iz*i for i in range ( iz + 1 ) ] )
ix , iy , iz  =  8 ,  8 ,  8
mc_xy = h2_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                  [ ymax/iy*i for i in range ( iy + 1 ) ] )
ix , iy , iz  =  9 ,  9 ,  9
mc_yz = h2_axes ( [ ymax/iy*i for i in range ( iy + 1 ) ] ,
                  [ zmax/iz*i for i in range ( iz + 1 ) ] )
ix , iy , iz  =  7 ,  7 ,  7
mc_xz = h2_axes ( [ xmax/ix*i for i in range ( ix + 1 ) ] ,
                  [ zmax/iz*i for i in range ( iz + 1 ) ] )
ix , iy , iz  = 95 , 95 , 95
mc_x  = h1_axis ( [ xmax/ix*i for i in range ( ix + 1 ) ] )
mc_y  = h1_axis ( [ ymax/iy*i for i in range ( iy + 1 ) ] )
mc_z  = h1_axis ( [ zmax/iz*i for i in range ( iz + 1 ) ] )

## prepare re-weighting machinery
maxIter = 15

## check database
if not os.path.exists( dbname ) :
    logger.info('Create new weights DBASE')
    with DBASE.open ( dbname , 'c' ) : ##  create new empty db
        pass
else :
    logger.info('Existing weights DBASE will be used')
    with DBASE.open ( dbname , 'r' ) as db :
        db.ls()

# =============================================================================
## make reweighting iterations

from   ostap.tools.reweight           import Weight, makeWeights,  WeightingPlot, W2Data
from   ostap.fitting.pyselectors      import Variable
import ostap.parallel.parallel_fill

# =============================================================================
## configuration of reweighting
weightings = (
    ## variable          address in DB
    Weight.Var (  'x'          , 'x-reweight'  ) ,
    Weight.Var (  'y'          , 'y-reweight'  ) ,
    Weight.Var (  'z'          , 'z-reweight'  ) ,
    Weight.Var ( ('x','y')     , 'XY-reweight' ) ,
    Weight.Var ( ('y','z')     , 'YZ-reweight' ) ,
    Weight.Var ( ('x','z')     , 'XZ-reweight' ) ,
    Weight.Var ( ('x','y','z') , '3D-reweight' ) ,
    )

# =============================================================================
## variables to be used in MC-dataset
variables  = [
    Variable ( 'x' , 'x-var'  , 0  , xmax ) ,
    Variable ( 'y' , 'y-var'  , 0  , ymax ) ,
    Variable ( 'z' , 'z-var'  , 0  , zmax ) ,
    ]

# =============================================================================
datatree   = ROOT.TChain ( tag_data ) ; datatree.Add ( testdata )
title = 'Data/target dataset'
logger.info ( '%s:\n%s' % ( title , datatree.table2 ( variables = [ 'x' , 'y' , 'z' ] ,
                                                      title     = title    ,
                                                      prefix    = '# '     ) ) )

vct_data  = datatree.statVct ( 'x,y,z' )
n_data    = len ( datatree )
## table of global statistics
glob_stat = [ ( '#' , 'Mahalanobis' , 'Hotelling' , 'KL/DATA-MC' , 'KL/MC-DATA' , 'KL-sym' ) ]
# =============================================================================
with timing ( 'Prepare initial MC-dataset:' , logger = logger ) :
    mctree   = ROOT.TChain ( tag_mc   ) ; mctree   .Add ( testdata )
    ## fill dataset from input MC tree
    mcds_ , _ = mctree.make_dataset ( variables = variables ,
                                      selection = '0<x && x<%s && 0<y && y<%s && 0<z && z<%s' % ( xmax , ymax, zmax ) )

# =============================================================================
## Configuration of reweighting plots
# =============================================================================
plots  = [
    WeightingPlot ( 'z:y:x' , 'weight' , '3D-reweight' , h3d_data , mc_3d ) ,
    WeightingPlot ( 'z:x'   , 'weight' , 'XZ-reweight' , hxz_data , mc_xz ) ,
    WeightingPlot ( 'z:y'   , 'weight' , 'YZ-reweight' , hyz_data , mc_yz ) ,
    WeightingPlot ( 'y:x'   , 'weight' , 'XY-reweight' , hxy_data , mc_xy ) ,
    WeightingPlot ( 'x'     , 'weight' , 'x-reweight'  , hx_data  , mc_x  ) ,
    WeightingPlot ( 'z'     , 'weight' , 'z-reweight'  , hz_data  , mc_z  ) ,
    WeightingPlot ( 'y'     , 'weight' , 'y-reweight'  , hy_data  , mc_y  ) ,
    ]

# =============================================================================
## start reweighting iterations:
for iter in range ( 1 , maxIter + 1 ) :

    tag = 'Reweighting iteration #%d' % iter
    logger.info ( allright ( tag ) )

    # =========================================================================
    ## 0) The weighter object
    weighter = Weight ( dbname , weightings )
    ## 1a) create new "weighted" mcdataset
    mcds = mcds_.Clone()

    with timing ( tag + ': add weight to MC-dataset' , logger = logger ) :
        # =========================================================================
        ## 1b) add  "weight" variable to dataset
        mcds.add_reweighting ( weighter ,  name = 'weight' , progress = True )
        title = 'Reweighted dataset at iteration #%d' % iter
        logger.info ( '%s:\n%s' % ( title , mcds.table2 ( variables = [ 'x' , 'y' , 'z' ] ,
                                                          title     = title    ,
                                                          cuts      = 'weight' ,
                                                          prefix    = '# '     ) ) )
        ## get MC data in a form of vector
        vct_mc = mcds.statVct ( 'x,y,z', 'weight' )
        n_mc   = int ( mcds.nEff('weight') )
        trow = ( '%d'   % iter ,
                 '%.4g' % vct_mc  .mahalanobis                 ( vct_data ) ,
                 '%.4g' % Ostap.Math.hotelling ( vct_mc , n_mc , vct_data , n_data ) ,
                 '%.4g' % vct_mc  .asymmetric_kullback_leibler ( vct_data ) ,
                 '%.4g' % vct_data.asymmetric_kullback_leibler ( vct_mc   ) ,
                 '%.4g' % vct_data.           kullback_leibler ( vct_mc   ) )
        glob_stat.append ( trow )
        ## title = 'MC-data as vector at #%s' % iter
        ## logger.info ( '%s\n%s' % ( title , vct_mc.table ( title = title , prefix = '# ' , correlations =True ) ) )

    # =========================================================================
    ## 2) update weights

    pmax  = 1.0 if iter <= 4 else 0.66 + 0.34 * math.exp ( - ( iter - 4.0 ) / 4.0 )
    power = lambda n : pmax if n <= 1 else pmax * ( 0.66666 / n + 0.33333 )

    if   iter < 3 : the_plots = plots [ : 1 ]
    elif iter < 4 : the_plots = plots [ : 4 ]
    else          : the_plots = plots

    ## weigth truncation: avoid very large change of weights for  single iteration
    wtruncate = ( 0.1 , 10 ) if iter < 4 else ( 0.5 , 2.0 )

    with timing ( tag + ': make actual reweighting:' , logger = logger ) :

        # =========================================================================
        ## 2a) the most important line: perform single iteration step
        active , _ = makeWeights (
            mcds                   , ## what to be reweighted
            the_plots              , ## reweighting plots/setup
            dbname                 , ## DBASE with reweighting constant
            delta      = 0.02      , ## stopping criteria
            minmax     = 0.08      , ## stopping criteria
            power      = power     , ## tune: effective power
            wtruncate  = wtruncate , ## truncate weights
            make_plots = True      , ## make control plots
            tag        = tag       ) ## tag for printout

    if not active and 5 <= iter :
        logger.info    ( allright ( 'No more iterations, converged after #%d' % iter ) )
        break

    mcds.clear()
    del mcds

else :

    logger.error ( "No convergency!" )

# ===========================================================================
title = 'Weighter object'
logger.info ( '%s:\n%s' % ( title , weighter.table ( prefix = '# ' ) ) )
# ============================================================================
## draw the convergency graphs
graphs = weighter.graphs ()
for key in graphs :
    with use_canvas ( "Convergency graph for '%s'" % key ) :
        graph = graphs [ key ]
        graph.draw ( 'a' )
# =============================================================================

# =============================================================================
with timing ( "Add weight column to initial MC-tree" , logger = logger ) :
    mctree   = ROOT.TChain ( tag_mc   ) ; mctree   .Add ( testdata )
    weighter = Weight ( dbname , weightings )
    mctree   = mctree.add_reweighting ( weighter ,  name = 'weight' )
    mctree   = ROOT.TChain ( tag_mc   ) ; mctree   .Add ( testdata )

# =============================================================================
## compare DATA  and MC before and after reweighting
# =============================================================================

datatree   = ROOT.TChain ( tag_data ) ; datatree.Add ( testdata )
mctree     = ROOT.TChain ( tag_mc   ) ; mctree  .Add ( testdata )
# =============================================================================
title = 'Data/target dataset'
logger.info ( '%s:\n%s' % ( title , datatree.table2 ( variables = [ 'x' , 'y' , 'z' ] ,
                                                      title     = title    ,
                                                      prefix    = '# '     ) ) )
# =============================================================================
title = 'MC-tree before reweighting'
logger.info ( '%s:\n%s' % ( title , mctree.table2   ( variables = [ 'x' , 'y' , 'z' ] ,
                                                      title     = title    ,
                                                      prefix    = '# '     ) ) )
# =============================================================================
title = 'MC-tree after reweighting'
logger.info ( '%s:\n%s' % ( title , mctree.table2   ( variables = [ 'x' , 'y' , 'z' ] ,
                                                      title     = title    ,
                                                      cuts      = 'weight' ,
                                                      prefix    = '# '     ) ) )

# =============================================================================
vct_final = mctree.statVct ( 'x,y,z' , 'weight' )
n_mc = int ( mctree.nEff('weight')  )
trow = ( '*' ,
         '%.4g'  % vct_final .mahalanobis                 ( vct_data  ) ,
         '%.4g'  % Ostap.Math.hotelling ( vct_final , n_mc ,  vct_data  , n_data ) ,
         '%+.4g' % vct_final .asymmetric_kullback_leibler ( vct_data  ) ,
         '%+.4g' % vct_data  .asymmetric_kullback_leibler ( vct_final ) ,
         '%+.4g' % vct_data  .           kullback_leibler ( vct_final ) )
glob_stat.append ( trow )

title = 'Global DATA/MC similarity'
table = T.table ( glob_stat , title = title , prefix = '# ' )
logger.info ( '%s\n%s' % ( title , table ) )

# =============================================================================
##                                                                      The END
# =============================================================================