import os
import contextlib
from PyPDF2 import PdfFileMerger
from pathlib import Path
import pandas as pd
import telebot
import requests
import resource
from sqlalchemy import create_engine
import time

from os import getenv
from dotenv import load_dotenv
load_dotenv()
PG_BF_USER_ENGINE = getenv('PG_BF_USER_ENGINE')
TELE_TOKEN_OZON_BOT = getenv('TELE_TOKEN_OZON_BOT')
OZON_API_CID = getenv('OZON_API_CID')
OZON_API_KEY = getenv('OZON_API_KEY')

bot = telebot.TeleBot(TELE_TOKEN_OZON_BOT)
engine = create_engine(PG_BF_USER_ENGINE, echo=True)

@bot.message_handler(content_types=['document'])
def docs_functions(message):

    resource.setrlimit(resource.RLIMIT_NOFILE, (65536, 65536))  # increase the limit of open files at one time in ubuntu

    # Генерация стикеров с таблицей по отправленному в чат csv файлу из лк озона
    if message.document.file_name == 'ozon.csv':

        # получение и сохранение таблицы из чата
        file_info = bot.get_file(message.document.file_id)
        csv_file = requests.get(f'https://api.telegram.org/file/bot{TELE_TOKEN_OZON_BOT}/{file_info.file_path}')
        with open('/home/analyst1/www/ozon_stickers/ozon_lk.csv', 'wb') as file:
            file.write(csv_file.content)

        bot.reply_to(message, f'Таблица создается...')

        # загрузка сохраненной таблицы с сервера
        input_csv = pd.read_csv('/home/analyst1/www/ozon_stickers/ozon_lk.csv', delimiter=';', dtype={'Артикул': str})
        input_csv.dropna(subset=['Номер заказа'], inplace=True)

        # загрузка 1с инфы из базы данных
        sql_query = '''
                        SELECT
                            vendor_code,
                            barcode,
                            brand,
                            item_name,
                            characteristic,
                            cell
                        FROM public.products_info_1c
                        '''

        prod_info_1c = pd.read_sql(sql_query, engine)

        # удаление ненужных данных из озоновской таблицы и форматирование
        orders_fbs_postings = input_csv[['Номер отправления', 'Артикул', 'Количество']]
        orders_fbs_postings.rename(columns={'Номер отправления': 'posting_number',
                                            'Артикул': 'vendor_code',
                                            'Количество': 'quantity'}, inplace=True)

        # мердж с данными из 1с
        orders_fbs_postings = orders_fbs_postings.merge(prod_info_1c, how='left', on='vendor_code')

        # сортировка для склада
        posting_sku_count = orders_fbs_postings.groupby(['posting_number'], as_index=False).agg(
            {'vendor_code': 'count'}).sort_values('vendor_code', ascending=True)
        posting_sku_count_1 = posting_sku_count[posting_sku_count['vendor_code'] == 1]
        posting_sku_count_more_than_1 = posting_sku_count[posting_sku_count['vendor_code'] > 1]

        orders_fbs_postings_1 = orders_fbs_postings[
            orders_fbs_postings.posting_number.isin(list(posting_sku_count_1.posting_number))].sort_values('item_name')
        orders_fbs_postings_more_than_1 = orders_fbs_postings[
            orders_fbs_postings.posting_number.isin(list(posting_sku_count_more_than_1.posting_number))].sort_values(
                                                                                ['posting_number',
                                                                                'item_name'])
        orders_fbs_postings = pd.concat([orders_fbs_postings_1, orders_fbs_postings_more_than_1])

        # выгрузка эксельки
        orders_fbs_postings = orders_fbs_postings[
            ['posting_number', 'item_name', 'characteristic', 'quantity', 'vendor_code', 'cell', 'barcode']]
        orders_fbs_postings.to_excel('/home/analyst1/www/ozon_stickers/ozon_stickers_table.xlsx', index=False)

        bot.reply_to(message, f'Стикеры создаются...')

        # запрос в апи озона не генерацию стикеров по списку номеров отправлений
        posting_numbers_list = list(orders_fbs_postings.posting_number.drop_duplicates())

        headers = {
            'Host': 'api-seller.ozon.ru',
            'Client-Id': OZON_API_CID,
            'Api-Key': OZON_API_KEY,
            'Content-Type': 'application/json'
        }

        url = 'https://api-seller.ozon.ru/v2/posting/fbs/package-label'

        # генерация пдфок по каждому отправлению и их сохранение в папку
        postings_quantity = len(posting_numbers_list)
        for i in range(postings_quantity):

            if i == round(postings_quantity * 0.1):
                bot.reply_to(message, f'Готово 10%...')
            if i == round(postings_quantity * 0.2):
                bot.reply_to(message, f'Готово 20%...')
            if i == round(postings_quantity * 0.3):
                bot.reply_to(message, f'Готово 30%...')
            if i == round(postings_quantity * 0.4):
                bot.reply_to(message, f'Готово 40%...')
            if i == round(postings_quantity * 0.5):
                bot.reply_to(message, f'Готово 50%...')
            if i == round(postings_quantity * 0.6):
                bot.reply_to(message, f'Готово 60%...')
            if i == round(postings_quantity * 0.7):
                bot.reply_to(message, f'Готово 70%...')
            if i == round(postings_quantity * 0.8):
                bot.reply_to(message, f'Готово 80%...')
            if i == round(postings_quantity * 0.9):
                bot.reply_to(message, f'Готово 90%...')

            query = {"posting_number": [posting_numbers_list[i]]}

            r = requests.post(url, headers=headers, json=query)

            if 'INVALID_ARGUMENT' in r.text:
                continue

            while r.text.startswith('%PDF') is False:
                time.sleep(1)
                r = requests.post(url, headers=headers, json=query)

            file_bytes = r.content

            with open(f'/home/analyst1/www/ozon_stickers/temp/sticker-{i:04d}.pdf', 'wb') as file:
                file.write(file_bytes)

        # объединение первых страниц всех пдфок в одну и её сохранение
        list_of_pdfs = os.listdir('/home/analyst1/www/ozon_stickers/temp/')
        list_of_pdfs.sort()

        with contextlib.ExitStack() as stack:
            merger = PdfFileMerger()
            files = [stack.enter_context(open(f'/home/analyst1/www/ozon_stickers/temp/{pdf}', 'rb')) for pdf in list_of_pdfs]

            for f in files:
                merger.append(fileobj=f, pages=(0, 1))
            with open('/home/analyst1/www/ozon_stickers/ozon_stickers.pdf', 'wb') as f:
                merger.write(f)

        # очистка папки temp
        [f.unlink() for f in Path('/home/analyst1/www/ozon_stickers/temp/').glob('*') if f.is_file()]

        # отправка результата в чат
        with open('/home/analyst1/www/ozon_stickers/ozon_stickers_table.xlsx', 'rb') as doc_xlsx:
            bot.send_document(message.chat.id, doc_xlsx)

        with open('/home/analyst1/www/ozon_stickers/ozon_stickers.pdf', 'rb') as doc_pdf:
            bot.send_document(message.chat.id, doc_pdf)

    elif '.csv' not in message.document.file_name:
        bot.reply_to(message, 'Неверный формат файла. Допустимый формат: .csv.')
    else:
        bot.reply_to(message, 'Неверное имя файла, допустимые имена: ozon.csv')
        # bot.reply_to(message, '1) "ozon.csv" - для получения стикеров')
        # bot.reply_to(message, '2) "заказы" - для получения артикула товара по номеру заказа WB')

@bot.message_handler(commands=['status'])
def show_status(message):

    bot.reply_to(message, f'Бот онлайн')

@bot.message_handler(commands=['info'])
def show_status(message):

    bot.reply_to(message, 'Для получения стикеров нужно отправить таблицу csv из лк озона.\n '
                          'Название таблицы должно быть "ozon.csv"\n '
                          'В таблице должны быть колонки: "Номер отправления", "Артикул", "Количество"')

bot.infinity_polling()


from PyQt5.QtWidgets import QWidget, QLabel, QVBoxLayout, QScrollArea, QGridLayout, QLineEdit, QPushButton, QButtonGroup, QMessageBox
from PyQt5.QtWidgets import QHBoxLayout, QDialog, QDialogButtonBox

class TableListView:
    def __init__(self, window, root_window):
        self.main_widget = window
        self.root_window = root_window

        self.drawToolBar()

    def __del__(self):
        self.main_widget_draw_area.deleteLater()
        self.scrollable.deleteLater()
        self.scrollable_v_lay.deleteLater()
        self.main_widget.deleteLater()
        self.button_group_toolbar.deleteLater()

    def initDraw(self, width, height):
        self.main_widget.setMinimumWidth(width)
        self.main_widget.setMinimumHeight(height)

        self.scrollable = QScrollArea(self.main_widget)
        self.scrollable.move(0, 70)
        self.scrollable_v_lay = QVBoxLayout()
        self.scrollable.setLayout(self.scrollable_v_lay)
        self.main_widget_draw_area = QWidget(self.main_widget)
        self.scrollable_v_lay.addWidget(self.main_widget_draw_area)

        self.main_widget.show()
        pass

    def drawToolBar(self):
        self.button_group_toolbar = QButtonGroup(self.main_widget)

        self.search_b_g_tb = QPushButton(self.main_widget)
        self.search_b_g_tb.setText("Konsola SQL")

        self.button_group_toolbar.addButton(self.search_b_g_tb)


from PyQt5 import QtCore, QtWidgets, QtGui

from Item import Item
import Constant as c
from Tikzifyables.Labelable import Labelable
from Tikzifyables.DashPatternable import DashPatternable
from Tikzifyables.Colourable.LineColourable import LineColourable
from Tikzifyables.Colourable.FillColourable import FillColourable

class Point(Item, Labelable, DashPatternable, LineColourable, FillColourable):
    def __init__(self, item):
        """Construct Point."""
        Item.__init__(self, item)
        if item is None:
            self.dictionary_builder(None, "", "")
        Labelable.__init__(self, self.item)
        DashPatternable.__init__(self, self.item)
        LineColourable.__init__(self, self.item)
        FillColourable.__init__(self, self.item)
        self.__canvas_x = 0
        self.__canvas_y = 0

    def get_canvas_coordinates(self):
        return self.__canvas_x, self.__canvas_y

    def set_canvas_coordinates(self, x, y):
        self.__canvas_x = x
        self.__canvas_y = y

    def __tikzify_marker(self):
        extra = ''
        if self.item["marker"]["shape"] == 'regular polygon':
            extra = 'regular polygon sides=' + str(self.item["marker"]["shape_number"])
        elif self.item["marker"]["shape"] == 'star':
            extra = f'star points={self.item["marker"]["shape_number"]}, star point ratio={self.item["marker"]["ratio"]}'

        options = [
            self.item["marker"]["shape"],
            'minimum size=' + str(self.item["marker"]["size"]),
            'inner sep=' + str(self.item["marker"]["inner_sep"]),
            '' if self.item["marker"]["shape"] == 'circle' or self.item["marker"]["rounded_corners"] == c.Point.Default.Marker.ROUNDED_CORNERS else f'rounded corners={self.item["marker"]["rounded_corners"]}pt',
            '' if self.item["marker"]["text_width"] == c.Point.Default.Marker.TEXT_WIDTH else 'text width=' + str(self.item["marker"]["text_width"]),
        ]
        if extra:
            options.append(extra)
        options = filter(bool, options)

        return ', '.join(options)

    def tikzify_node(self):
        options = [
            f'draw={c.Point.Default.Line_Colour.NAME}' if not bool(self.tikzify_line_colour()) else f'draw={self.tikzify_line_colour()}',
            f'fill={c.Point.Default.FillColor.NAME}' if not bool(self.tikzify_fill_colour()) else f'fill={self.tikzify_fill_colour()}',
            '' if self.item["line"]["line_width"] == c.Point.Default.LINE_WIDTH else f'line width={self.item["line"]["line_width"]}',
            self.__tikzify_marker(),
            self.tikzify_dash()
        ]
        options = filter(bool, options)
        return '\\node[%s] (%s) at (%s) {%s};' % (', '.join(options), self.get_id(), self.get_id(), self.item["marker"]["text"])

    def tikzify(self, items=None):
        raise NotImplementedError

    def draw_on_canvas(self, items, scene, colour=QtCore.Qt.darkMagenta):
        radius = 6
        graphics_point = QtWidgets.QGraphicsEllipseItem(self.__canvas_x-radius,self.__canvas_y-radius, 2*radius, 2*radius)
        graphics_point.setBrush(QtGui.QBrush(colour))
        scene.addItem(graphics_point)
        if scene.show_canvas_labels:
            text = QtWidgets.QGraphicsTextItem(self.item["id"])
            text.setPos(self.__canvas_x,self.__canvas_y)
            text.setDefaultTextColor (colour)
            font = text.font()
            font.setPointSize(12)
            font.setWeight(600)
            text.setFont(font)
            # text.setTextWidth(400)
            scene.addItem(text)

    def distance_sqr(self, x, y, items):
        return (self.__canvas_x-x) ** 2 + (self.__canvas_y-y) ** 2

    def next_id_func(self, definition, iter_counter):
        return chr(ord('A') + iter_counter % 26) + (iter_counter // 26) * '\''

    def dictionary_builder(self, definition, id_, sub_type=None):
        dictionary = {}
        dictionary["id"] = id_
        dictionary["type"] = 'point'
        dictionary["sub_type"] = sub_type
        dictionary["show"] = True
        dictionary["definition"] = definition
        dictionary["marker"] = {}
        dictionary["marker"]["size"] = c.Point.Default.Marker.SIZE
        dictionary["marker"]["shape"] = c.Point.Default.Marker.SHAPE
        dictionary["marker"]["shape_number"] = c.Point.Default.Marker.SHAPE_NUMBER
        dictionary["marker"]["ratio"] = c.Point.Default.Marker.RATIO
        dictionary["marker"]["inner_sep"] = c.Point.Default.Marker.INNER_SEP
        dictionary["marker"]["rounded_corners"] = c.Point.Default.Marker.ROUNDED_CORNERS
        dictionary["marker"]["text"] = c.Point.Default.Marker.TEXT
        dictionary["marker"]["text_width"] = c.Point.Default.Marker.TEXT_WIDTH
        dictionary["line"] = {}
        dictionary["line"]["line_width"] = c.Point.Default.LINE_WIDTH
        dictionary["line"]["colour"] = {}
        dictionary["line"]["colour"]["name"] = c.Point.Default.Line_Colour.NAME
        dictionary["line"]["colour"]["mix_with"] = c.Point.Default.Line_Colour.MIX_WITH
        dictionary["line"]["colour"]["mix_percent"] = c.Point.Default.Line_Colour.MIX_RATIO
        dictionary["line"]["colour"]["strength"] = c.Point.Default.Line_Colour.STRENGTH
        dictionary["line"]["dash"] = {}
        dictionary["line"]["dash"]["stroke"] = c.Point.Default.LINE_DASH_STROKE
        dictionary["line"]["dash"]["custom_pattern"] = c.Point.Default.LINE_DASH_CUSTOM
        dictionary["fill"] = {}
        dictionary["fill"]["colour"] = {}
        dictionary["fill"]["colour"]["name"] = c.Point.Default.Fill_Colour.NAME
        dictionary["fill"]["colour"]["mix_with"] = c.Point.Default.Fill_Colour.MIX_WITH
        dictionary["fill"]["colour"]["mix_percent"] = c.Point.Default.Fill_Colour.MIX_RATIO
        dictionary["fill"]["colour"]["strength"] = c.Point.Default.Fill_Colour.STRENGTH
        dictionary["label"] = {}
        dictionary["label"]["show"] = c.Point.Default.Label.SHOW
        dictionary["label"]["text"] = f'${id_}$'
        dictionary["label"]["text"] = c.Point.Default.Label.TEXT
        dictionary["label"]["anchor"] = c.Point.Default.Label.ANCHOR
        dictionary["label"]["offset"] = c.Point.Default.Label.OFFSET

        self.item = dictionary

#


'''
Created on 15 Aug 2011

@author: samgeen
'''

' TODO: CREATE POLYMORPHIC INTERFACE!!! '

import pyglet
from pyglet.gl import *
from pyglet.gl.glu import *

import numpy as np

#import EventHandler as Events
#from MouseHandler import MouseHandler
import math

from .Maths.Vect3 import Vect3
from .Maths.Quaternion import Rotation

from . import Events

class CameraController(object):
    def __init__(self, camera, window):
        self._lmb = False
        self._camera = camera
        self._window = window
        self._userInput = Events.UserInputHandler(self)
        self._window.push_handlers(self._userInput)

    def OnMouseMove(self, state):
        '''
        Parse mouse movement
        '''
        if self._lmb:
            self._camera.OnMouseMove(state)
        else:
            self._camera.ResetPos(state)

    def OnKeyboard(self, state):
        '''
        Register the pressing or releasing of a keyboard key
        This can be overridden by the concrete class if desired, or left inactive
        state - a dictionary of the mouse state:
                   {"button": button, "mod": modifier keys used, "pressed":True or False}
        '''
        # Escape
        pressed = state["pressed"]
        button = state["button"]

        # Q
        if button == pyglet.window.key.Q:
            self._camera.ZoomIn(pressed)
        # E
        if button == pyglet.window.key.E:
            self._camera.ZoomOut(pressed)

    def OnMouseButton(self, state):
        if state["button"] % 2 == pyglet.window.mouse.LEFT:
            self._lmb = state["pressed"] # True if pressed, False if not

class Camera(object):
    '''
    classdocs
    '''

    def __init__(self, window):
        '''
        Constructor
        '''
        zoomInit = 2.0
        self.__zoom = zoomInit
        self.__lookAt = Vect3(0.0,0.0,0.0)
        self.__position = Vect3(0.0,0.0,zoomInit)
        self.__up = Vect3(0.0,1.0,0.0)
        self.__right = Vect3(1.0,0.0,0.0) # This is reset during the view calculations
        self.__xangle = 0.0
        self.__yangle = 0.0
        self.__zoomSpeed = 1.30
        # Keep track of time for smooth moving of camera
        self._clock = pyglet.clock.Clock()
        self._time = self._clock.time()
        self._prevtime = self._time
        self._zoomIn = False
        self._zoomOut = False
        # Keyboard and mouse handler
        self._controller = CameraController(self, window)
        # PREVIOUS MOUSE POSITIONS
        self.__oldscrx = 0
        self.__oldscry = 0

    def ResetPos(self, state):
        '''
        Reset position when the user begins scrolling again
        '''
        x,y = state["pos"]
        self.__oldscrx = x
        self.__oldscry = y

    def OnMouseMove(self, state):
        '''
        Parse mouse movement
        '''
        x, y = state["pos"]
        # Good speed for a mouse
        rotateSpeed = 5e-2
        # Good speed for a wiimote
        rotateSpeed = 2e-2
        xangle = -(x - self.__oldscrx) * rotateSpeed
        yangle = (y - self.__oldscry) * rotateSpeed
        self.__oldscrx = x
        self.__oldscry = y
        self.Rotate((xangle, yangle))

    def Rotate(self, angles):
        self.__xangle += angles[0]
        self.__yangle += angles[1]
        # Find total angle
        angle = math.sqrt(angles[0]*angles[0] + angles[1]*angles[1])
        # Make sure the right-pointing vector is set properly (as cross product of up and lookAt)
        self.__right = -self.__up ^ self.__position.Normal()
        # Find direction the angle displacement is pointing
        displacement = self.__right*angles[0] + self.__up*angles[1]
        # Now find the vector perpendicular to that and the camera direction
        vector = -self.__position ^ displacement
        # Now ROTATE
        if vector.LengthSqrd() > 0:
            rotation = Rotation(angle, vector)
            self.__position = rotation.RotateVector(self.__position)
            self.__up = rotation.RotateVector(self.__up)
        # Done!
        Events.Redraw()

    def Zoom(self, zoom=None):
        '''
        Set zoom or return its current value (or both)
        '''
        if zoom:
            self.__position *= zoom / self.__zoom
            self.__zoom = zoom
        return self.__zoom

    def ZoomActive(self):
        '''
        HACK - Is the zoom currently active?
        '''
        return self._zoomIn or self._zoomOut

    def ZoomIn(self, on=True):
        '''
        Zoom in camera
        '''
        self._zoomIn = on

    def ZoomOut(self, on=True):
        '''
        Zoom in camera
        '''
        self._zoomOut = on

    def Draw(self):
        '''
        Update the display with the current camera view
        '''
        self._prevtime = self._time
        self._time = self._clock.time()
        dt = self._time - self._prevtime
        self._CalcZoom(dt)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45,1.0,0.01,10.0)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt(self.__position[0], self.__position[1], self.__position[2],
                  self.__lookAt[0], self.__lookAt[1], self.__lookAt[2],
                  self.__up[0], self.__up[1], self.__up[2])

    def LookAt(self, point):
        '''
        Set point to look at
        '''
        self.__lookAt = point

    def Position(self, point=0):
        '''
        Set/get position of camera
        '''
        if type(point) == type(0):
            return self.__position
        else:
            self.__position = point

    def Up(self, direction):
        '''
        Give direction of up
        '''
        self.__up = direction

    def _CalcZoom(self, dt):
        # The zoom is a logarithmic zoom, i.e. zooms by factors of 10 in constant time
        zoomfact = self.__zoomSpeed ** dt
        if self._zoomIn:
            self.__zoom /= zoomfact
            self.__position /= zoomfact
        if self._zoomOut:
            self.__zoom *= zoomfact
            self.__position *= zoomfact


# Should be considered a singleton representing a currently loaded project.

import json
import csv
import re
import os
import shutil
from collections import namedtuple, defaultdict
from main_utils import appPath
from dotmap import DotMap

import interop.jplag_preprocessor as jplag_preprocessor
import interop.jplag_runner as jplag_runner
import interop.moodle_downloader as moodle_downloader

_CurrentProjectPath = None  # initially no project file is loaded
_CurrentAssignment = None  # no detection performed

def _templatesDir():
    return appPath()  # may be revised

def settings():
    assert _CurrentProjectPath
    json_path = os.path.join(_CurrentProjectPath, "config.json")
    with open(json_path) as f:
        return DotMap(json.load(f))

def saveSettings(config):
    json_path = os.path.join(_CurrentProjectPath, "config.json")
    with open(json_path, "w") as f:
        json.dump(config, f, indent=4)

def newProject(project_path):  # initialize project at given path
    global _CurrentProjectPath
    _CurrentProjectPath = project_path
    shutil.copyfile(
        os.path.join(_templatesDir(), "config_template.json"),
        os.path.join(_CurrentProjectPath, "config.json"),
    )

def openProject(project_path):  # open project at given path
    global _CurrentProjectPath
    _CurrentProjectPath = project_path

def syncWithDataSource(progress_object):
    assert _CurrentProjectPath
    os.chdir(_CurrentProjectPath)
    moodle_downloader.download(settings(), progress_object)

def assignments():
    assert _CurrentProjectPath
    os.chdir(os.path.join(_CurrentProjectPath, settings().moodle_submissions_dir))
    r = [d for d in os.listdir() if os.path.isdir(d)]
    os.chdir(_CurrentProjectPath)
    return r

def _is_arc_student(student):
    return student.startswith("arc[")

def _filter_arc_records(log):
    def student_name(student_name):
        m = re.match("(arc\[.+\]).+", student_name)
        return m.group(1) if m else student_name

    # log is in a format
    # "<user1>-<user2>: <sim_ratio>\n<user1>-<user3>: <sim_ratio>\n..."
    # archived users are "arc[dirname]user"

    Match = namedtuple("Match", ["student1", "student2", "score"])

    log_matches = []  # first, remove all archive-archive matches
    for log_line in log.strip().split("\n"):
        m = re.match("(.+)-(.+): (.+)", log_line)
        log_matches.append(Match(m.group(1), m.group(2), float(m.group(3))))

    # next, combine all students of the same archive
    # use max score over the archive
    log_filtered = defaultdict(float)
    for m in log_matches:
        if not (_is_arc_student(m.student1) and _is_arc_student(m.student2)):
            key = (student_name(m.student1), student_name(m.student2))

            if key not in log_filtered:  # name pair not found
                key = tuple(reversed(key))  # swap names

            log_filtered[key] = max(log_filtered[key], m.score)

    # TODO: use dictionary as an input format for the matrix
    return "\n".join([f"{n1}-{n2}: {v}" for (n1, n2), v in log_filtered.items()])

def detect(asgn):
    global _CurrentAssignment
    assert _CurrentProjectPath

    cfg = settings()
    os.chdir(_CurrentProjectPath)
    _CurrentAssignment = asgn

    # do not run if we already have results
    if not os.path.exists(f"jpl_out_{asgn}.log"):
        jplag_preprocessor.preprocess_dirs(
            cfg.moodle_submissions_dir,
            cfg.archive_dirs,
            cfg.assignment_regexes,
            asgn,
        )
        jplag_runner.run(
            cfg.java_path,
            cfg.template_dir,
            cfg.language,
            cfg.min_sim,
            cfg.jplag_args,
            asgn,
        )

    with open(f"jpl_out_{asgn}.log") as f:
        return _filter_arc_records(f.read())

def _find_archive_matches(student1, student2, asgn_path):
    # if one of the students is arc, we need to rename
    # it according to the best match from the log
    # there is some intersection with code above

    # here we simply have someone with archive prefix
    maxsim = 0
    r1 = r2 = ""
    with open(f"{asgn_path}.log") as f:
        for log_line in f.read().strip().split("\n"):
            m = re.match("(.+)-(.+): (.+)", log_line)
            name1, name2, sim = m.group(1), m.group(2), float(m.group(3))
            opt1 = name1.startswith(student1) and name2.startswith(student2)
            opt2 = name2.startswith(student1) and name1.startswith(student2)
            if sim > maxsim and (opt1 or opt2):
                maxsim, r1, r2 = sim, name1, name2

    return r1, r2

def htmlReportPath(student1, student2):
    assert _CurrentProjectPath
    assert _CurrentAssignment

    asgn_path = os.path.join(_CurrentProjectPath, f"jpl_out_{_CurrentAssignment}")
    csv_path = os.path.join(asgn_path, "pair_report.csv")

    if _is_arc_student(student1) or _is_arc_student(student2):
        student1, student2 = _find_archive_matches(student1, student2, asgn_path)

    # student1; student2; html-report
    with open(csv_path) as f:
        reader = csv.reader(f, delimiter=";")
        for row in reader:
            if (student1, student2) in [(row[0], row[1]), (row[1], row[0])]:
                return os.path.join(asgn_path, row[2])

    assert False, "match not found"

#### MP ####
# NOTE(mm): obsolete function, but referenced in one place
def viewMatchReport(studentID_1, studentID_2):
    # TODO: should be in main window
    pass

    # # TODO: always shows 'match0.html' -- for testing
    # html_path = os.path.join(
    #     _CurrentProjectPath, f"jpl_out_{_CurrentAssignment}", "match0.html"
    # )
    # isOk = MatchViewerDialog.show(_getMainWin(), studentID_1, studentID_2, html_path)
    # if isOk:
    #     print("Yay!")
    # CurrentAssignment
    # do whatever you want

def dataSource():
    # Здесь, по моему замыслу, настраивается доступ к Мудлу или его альтернативам.
    # Я могу сделать удобную форму для этого.
    # Идея в том, что если у тебя уже есть готовый проект, но что-то поменялось,
    # например, пароль к Мудлу или еще что-то,
    # ты можешь зайти сюда и поменять настройки.
    print("data source")

def importExportSettings():
    # Эта функция на случай, если ты уже все под себя настроил в одном проекте,
    # и не хочешь заниматься тем же самым в новом.
    # Думаю, можно просто дать возможность выбрать файл настроек из другого проекта,
    # и скопировать их в текущий. Понятно, что это не самая жизненно необходимая функция,
    # но очень может быть полезной, а сделать несложно… наверное.
    print("import and export settings")

def exportTemplate():
    # Да, здесь я признаю, увлеклась.
    # Даже точно не помню какой изначально был план :)
    # Наверное, это будет зависеть от того, что такое будет наш проект,
    # и что будет создаваться вначале. Но, вообще,
    # было бы неплохо иметь возможность загрузить некий каркас,
    # возможно с примерами, чтобы пользователь мог попробовать все, до того,
    # как переименует все файлы нужным образом, и еще так по-всякому по-разному заморозится.
    # Но, это, конечно та функция, которую ты ужалишь первой, и на том и будет, наш компромисс :D
    # Как-то так...
    print("export template")


"""Posts models configuration"""
from ckeditor.fields import RichTextField
from django.contrib.auth import get_user_model
from django.db import models
from django.urls import reverse

User = get_user_model()

class Category(models.Model):
    """A category for post."""

    title = models.CharField(max_length=20)

    def __str__(self):
        return self.title

    class Meta:
        verbose_name_plural = "categories"

class Comment(models.Model):
    """A comment for post."""

    content = models.TextField()
    timestamp = models.DateTimeField(auto_now_add=True)
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    post = models.ForeignKey(
        "Post", related_name="comments", on_delete=models.CASCADE
    )

    def __str__(self):
        return self.user.username

STATUS = ((0, "Non publié"), (1, "Publié"))

class Post(models.Model):
    """A post for user."""

    title = models.CharField(max_length=100)
    author = models.ForeignKey(User, on_delete=models.CASCADE)
    categories = models.ManyToManyField(Category)
    updated_on = models.DateTimeField(auto_now=True)
    overview = models.CharField(max_length=255)
    content = RichTextField()
    created_on = models.DateTimeField(auto_now_add=True)
    thumbnail = models.ImageField(upload_to="images/")
    featured = models.BooleanField()
    status = models.IntegerField(choices=STATUS, default=0)

    def __str__(self):
        return self.title + " | " + str(self.author)

    def get_absolute_url(self):
        """get_absolute_url function allows to redirect to the home page."""
        return reverse("home")

    @property
    def get_comments(self):
        """
        get_comments is a function to get comments
        and return an object ordering by timestamp
        """
        return self.comments.all().order_by("-timestamp")

    @property
    def comment_count(self):
        """
        comment_count is a function to count comments
        per post and return an object
        """
        return Comment.objects.filter(post=self).count()


from functools import wraps

def log_function_data(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        """I'm wrapper function"""
        print(f"you are about to call {fn.__name__}")
        print(f"Here's the documentation: {fn.__doc__}")
        return fn(*args, **kwargs)
    return wrapper

@log_function_data
def add(x, y):
    """Adds two numbers together."""
    return x + y

print(add(10,30))

#It will show the correct information because of the @wraps(fn) in the wrapper function.
print(add.__doc__)
print(add.__name__)
help(add)


import service.globals as global_variables
import component.filesystem_item as filesystem_item
import dialogs.add_filesystem as add_filesystem
import pybookeeping.core.operation.filesystem as filesystem
import PySide.QtGui as QtGui
import PySide.QtCore as QtCore
import json
import functools

class FilesystemList(QtGui.QListWidget):
	_stylesheet = """
		QListWidget {
			font-family: "Lucida Grande", Verdana, Helvetica, Arial, sans-serif;
			font-size: 14px;
			background-color: #a6a6a6;
		}

		QListWidget::item {
			color: #a6a6a6;
			background-color: white;
			padding-top: 10px;
			padding-bottom: 10px;
			border-bottom: 1px solid black;
		}

		QListWidget::item:selected {
			color: white;
			background-color: #5abae1;
		}

		QListWidget::item:focus {
			border: 0px;
		}
	"""

	def __init__(self, entity):
		QtGui.QListWidget.__init__(self, entity)
		self.setSortingEnabled(True)
		self.setStyleSheet(self._stylesheet)
		self.setFocusPolicy(QtCore.Qt.NoFocus)

		global_variables.twig_signal.add_filesystem.connect(self.create_new)
		global_variables.twig_signal.view_changed.connect(self.item_changed)
		self.itemSelectionChanged.connect(self.item_changed)

		self.filesystem = filesystem.Filesystem(global_variables.bookeeping_connection)
		self.display_children()

	def display_children(self):
		for filesystem_info in self.filesystem.get_all_filesystem(global_variables.userid)[1]:
			item_tooltip = json.dumps(filesystem_info, indent = 8)
			item = filesystem_item.FilesytemItem(filesystem_info, self)
			item.setToolTip(item_tooltip)

	def remove_children(self):
		while self.count() > 0:
			self.takeItem(0)

	def contextMenuEvent(self, event):
		menu = QtGui.QMenu()

		item = self.itemAt(event.x(), event.y())
		if item is None:
			menu.addAction(QtGui.QAction("Add Filesystem", menu, triggered = self.create_new))
		else:
			menu.addAction(QtGui.QAction("Show Changes", menu, triggered = functools.partial(item.filesystem_changes)))
			menu.addAction(QtGui.QAction("Commit", menu, triggered = functools.partial(item.commit_changes)))

		menu.exec_(QtGui.QCursor.pos())

	def create_new(self):
		add_dialog = add_filesystem.AddFilesystem()
		if add_dialog.exec_() == QtGui.QDialog.Accepted:
			self.remove_children()
			self.display_children()

	def item_changed(self):
		current_item = self.currentItem()
		if current_item is not None:
			current_filesytem_info = self.currentItem().filesystem_info
			global_variables.twig_signal.filesystem_list_changed.emit(current_filesytem_info)


import unittest

import numpy as np

from pgmpy.factors.discrete import DiscreteFactor
from pgmpy.models import JunctionTree
from pgmpy.tests import help_functions as hf

class TestJunctionTreeCreation(unittest.TestCase):
    def setUp(self):
        self.graph = JunctionTree()

    def test_add_single_node(self):
        self.graph.add_node(("a", "b"))
        self.assertListEqual(list(self.graph.nodes()), [("a", "b")])

    def test_add_single_node_raises_error(self):
        self.assertRaises(TypeError, self.graph.add_node, "a")

    def test_add_multiple_nodes(self):
        self.graph.add_nodes_from([("a", "b"), ("b", "c")])
        self.assertListEqual(
            hf.recursive_sorted(self.graph.nodes()), [["a", "b"], ["b", "c"]]
        )

    def test_add_single_edge(self):
        self.graph.add_edge(("a", "b"), ("b", "c"))
        self.assertListEqual(
            hf.recursive_sorted(self.graph.nodes()), [["a", "b"], ["b", "c"]]
        )
        self.assertListEqual(
            sorted([node for edge in self.graph.edges() for node in edge]),
            [("a", "b"), ("b", "c")],
        )

    def test_add_single_edge_raises_error(self):
        self.assertRaises(ValueError, self.graph.add_edge, ("a", "b"), ("c", "d"))

    def test_add_cyclic_path_raises_error(self):
        self.graph.add_edge(("a", "b"), ("b", "c"))
        self.graph.add_edge(("b", "c"), ("c", "d"))
        self.assertRaises(ValueError, self.graph.add_edge, ("c", "d"), ("a", "b"))

    def tearDown(self):
        del self.graph

class TestJunctionTreeMethods(unittest.TestCase):
    def setUp(self):
        self.factor1 = DiscreteFactor(["a", "b"], [2, 2], np.random.rand(4))
        self.factor2 = DiscreteFactor(["b", "c"], [2, 2], np.random.rand(4))
        self.factor3 = DiscreteFactor(["d", "e"], [2, 2], np.random.rand(4))
        self.factor4 = DiscreteFactor(["e", "f"], [2, 2], np.random.rand(4))
        self.factor5 = DiscreteFactor(["a", "b", "e"], [2, 2, 2], np.random.rand(8))

        self.graph1 = JunctionTree()
        self.graph1.add_edge(("a", "b"), ("b", "c"))
        self.graph1.add_factors(self.factor1, self.factor2)

        self.graph2 = JunctionTree()
        self.graph2.add_nodes_from([("a", "b"), ("b", "c"), ("d", "e")])
        self.graph2.add_edge(("a", "b"), ("b", "c"))
        self.graph2.add_factors(self.factor1, self.factor2, self.factor3)

        self.graph3 = JunctionTree()
        self.graph3.add_edges_from([(("a", "b"), ("b", "c")), (("d", "e"), ("e", "f"))])
        self.graph3.add_factors(self.factor1, self.factor2, self.factor3, self.factor4)

        self.graph4 = JunctionTree()
        self.graph4.add_edges_from(
            [
                (("a", "b", "e"), ("b", "c")),
                (("a", "b", "e"), ("e", "f")),
                (("d", "e"), ("e", "f")),
            ]
        )
        self.graph4.add_factors(self.factor5, self.factor2, self.factor3, self.factor4)

    def test_check_model(self):
        self.assertRaises(ValueError, self.graph2.check_model)
        self.assertRaises(ValueError, self.graph3.check_model)
        self.assertTrue(self.graph1.check_model())
        self.assertTrue(self.graph4.check_model())

    def test_states(self):
        self.assertDictEqual(
            self.graph4.states,
            {
                "a": [0, 1],
                "b": [0, 1],
                "e": [0, 1],
                "c": [0, 1],
                "d": [0, 1],
                "f": [0, 1],
            },
        )

    def tearDown(self):
        del self.factor1
        del self.factor2
        del self.factor3
        del self.factor4
        del self.factor5

        del self.graph1
        del self.graph2
        del self.graph3
        del self.graph4

class TestJunctionTreeCopy(unittest.TestCase):
    def setUp(self):
        self.graph = JunctionTree()

    def test_copy_with_nodes(self):
        self.graph.add_nodes_from([("a", "b", "c"), ("a", "b"), ("a", "c")])
        self.graph.add_edges_from(
            [(("a", "b", "c"), ("a", "b")), (("a", "b", "c"), ("a", "c"))]
        )
        graph_copy = self.graph.copy()

        self.graph.remove_edge(("a", "b", "c"), ("a", "c"))
        self.assertFalse(self.graph.has_edge(("a", "b", "c"), ("a", "c")))
        self.assertTrue(graph_copy.has_edge(("a", "b", "c"), ("a", "c")))

        self.graph.remove_node(("a", "c"))
        self.assertFalse(self.graph.has_node(("a", "c")))
        self.assertTrue(graph_copy.has_node(("a", "c")))

        self.graph.add_node(("c", "d"))
        self.assertTrue(self.graph.has_node(("c", "d")))
        self.assertFalse(graph_copy.has_node(("c", "d")))

    def test_copy_with_factors(self):
        self.graph.add_edges_from([[("a", "b"), ("b", "c")]])
        phi1 = DiscreteFactor(["a", "b"], [2, 2], np.random.rand(4))
        phi2 = DiscreteFactor(["b", "c"], [2, 2], np.random.rand(4))
        self.graph.add_factors(phi1, phi2)
        graph_copy = self.graph.copy()

        self.assertIsInstance(graph_copy, JunctionTree)
        self.assertIsNot(self.graph, graph_copy)
        self.assertEqual(
            hf.recursive_sorted(self.graph.nodes()),
            hf.recursive_sorted(graph_copy.nodes()),
        )
        self.assertEqual(
            hf.recursive_sorted(self.graph.edges()),
            hf.recursive_sorted(graph_copy.edges()),
        )
        self.assertTrue(graph_copy.check_model())
        self.assertEqual(self.graph.get_factors(), graph_copy.get_factors())

        self.graph.remove_factors(phi1, phi2)
        self.assertTrue(
            phi1 not in self.graph.factors and phi2 not in self.graph.factors
        )
        self.assertTrue(phi1 in graph_copy.factors and phi2 in graph_copy.factors)

        self.graph.add_factors(phi1, phi2)
        self.graph.factors[0] = DiscreteFactor(["a", "b"], [2, 2], np.random.rand(4))
        self.assertNotEqual(self.graph.get_factors()[0], graph_copy.get_factors()[0])
        self.assertNotEqual(self.graph.factors, graph_copy.factors)

    def test_copy_with_factorchanges(self):
        self.graph.add_edges_from([[("a", "b"), ("b", "c")]])
        phi1 = DiscreteFactor(["a", "b"], [2, 2], np.random.rand(4))
        phi2 = DiscreteFactor(["b", "c"], [2, 2], np.random.rand(4))
        self.graph.add_factors(phi1, phi2)
        graph_copy = self.graph.copy()

        self.graph.factors[0].reduce([("a", 0)])
        self.assertNotEqual(
            self.graph.factors[0].scope(), graph_copy.factors[0].scope()
        )
        self.assertNotEqual(self.graph, graph_copy)
        self.graph.factors[1].marginalize(["b"])
        self.assertNotEqual(
            self.graph.factors[1].scope(), graph_copy.factors[1].scope()
        )
        self.assertNotEqual(self.graph, graph_copy)

    def tearDown(self):
        del self.graph


import os
import uuid
from pprint import pprint

import flask
import spotipy
from flask import session, request, redirect, render_template, abort, jsonify
from flask_login import login_required, current_user

from data import db_session
from data.users import User
from utils.spotify import spotify_login_required, session_cache_path, SCOPE, get_followed_artists

blueprint = flask.Blueprint(
    'music',
    __name__,
    template_folder='templates'
)

@blueprint.route('/music', methods=['GET', 'POST'])
@login_required
@spotify_login_required
def music(spotify: spotipy.Spotify):
    db_sess = db_session.create_session()

    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    return render_template('music.html', **params)

@blueprint.route('/music/search/<data>')
@login_required
@spotify_login_required
def search_music(data, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'artists': spotify.search(q=data, type='artist', limit=6)['artists']['items'],
        'albums': spotify.search(q=data, type='album', limit=6)['albums']['items'],
        'tracks': spotify.search(q=data, type='track', limit=6)['tracks']['items'],

        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    return render_template('search.html', **params)

@blueprint.route('/music/artist/followed')
@login_required
@spotify_login_required
def followed_artists(spotify: spotipy.Spotify):
    db_sess = db_session.create_session()

    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
        'artists': sorted(get_followed_artists(spotify), key=lambda x: x['popularity'],
                          reverse=True)
    }
    return render_template('followed_artists.html', **params)

@blueprint.route('/music/artist/<id>')
@login_required
@spotify_login_required
def artist(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()

    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    # 4K7kAAsdwnIHjgRk28OyOi

    try:
        params['artist'] = spotify.artist(id)
        params['top_tracks'] = spotify.artist_top_tracks(id)['tracks'][:6:]
        params['albums'] = spotify.artist_albums(id, album_type='album', limit=7)['items']
        params['singles'] = spotify.artist_albums(id, album_type='single', limit=7)['items']
        params['appears_on'] = spotify.artist_albums(id, album_type='appears_on', limit=7)['items']
    except:
        return abort(404)

    pprint(params['artist'])

    return render_template('artist.html', **params)

@blueprint.route('/music/artist/<id>/albums')
@login_required
@spotify_login_required
def artist_albums(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    try:
        params['artist'] = spotify.artist(id)
        params['albums'] = spotify.artist_albums(id, album_type='album', limit=50)['items']
    except:
        return abort(404)

    return render_template('artist_albums.html', **params)

@blueprint.route('/music/artist/<id>/singles')
@login_required
@spotify_login_required
def artist_singles(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    try:
        params['artist'] = spotify.artist(id)
        params['singles'] = spotify.artist_albums(id, album_type='single', limit=50)['items']
    except:
        return abort(404)

    return render_template('artist_singles.html', **params)

@blueprint.route('/music/artist/<id>/appears_on')
@login_required
@spotify_login_required
def artist_appears_on(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    try:
        params['artist'] = spotify.artist(id)
        params['appears_on'] = spotify.artist_albums(id, album_type='appears_on', limit=50)['items']
    except:
        return abort(404)

    return render_template('artist_appears_on.html', **params)

@blueprint.route('/music/playlist/<id>')
@login_required
@spotify_login_required
def playlist(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    # id = 2duj5GBXwLRdjqV9hjXy4o

    try:
        params['playlist'] = spotify.playlist(id)
        dur = 0
        for track in params['playlist']['tracks']['items']:
            dur += float(track['track']['duration_ms']) / 1000 / 60
        params['duration'] = round(dur, 2)
    except Exception:
        return abort(404)

    return render_template('playlist.html', **params)

@blueprint.route('/music/track/<id>')
@login_required
@spotify_login_required
def track(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    # id = 77yYxfpXB64ktXPVdU9xcF

    try:
        params['music'] = spotify.track(id)
        params['playlist'] = spotify.album(params['music']['album']['id'])
        dur = 0
        for track in params['playlist']['tracks']['items']:
            dur += float(track['duration_ms']) / 1000 / 60
        params['duration'] = round(dur, 2)
    except:
        return abort(404)

    return render_template('track.html', **params)

@blueprint.route('/music/album/<id>')
@login_required
@spotify_login_required
def album(id, spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {
        'current_user': db_sess.query(User).get(current_user.id),
        'spotify': spotify,
    }

    # id = 6WLBApYwMaajoFlzITzb6P

    try:
        params['playlist'] = spotify.album(id)
        dur = 0
        for track in params['playlist']['tracks']['items']:
            dur += float(track['duration_ms']) / 1000 / 60
        params['duration'] = round(dur, 2)
        params['artists'] = [[(artist['name'], artist['id']) for artist in track['artists']] for
                             track in params['playlist']['tracks']['items']]
        #
        #
        # сделать отображение всех артистов
        # использовать loop.first или loop.end
        #
    except:
        return abort(404)

    return render_template('album.html', **params)

@blueprint.route('/music/track/top')
@login_required
@spotify_login_required
def top_tracks(spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {'tracks': spotify.current_user_top_tracks(time_range='short_term'),
              'current_user': db_sess.query(User).get(current_user.id),
              'spotify': spotify,
              }

    return render_template('top_tracks.html', **params)

@blueprint.route('/music/artist/top')
@login_required
@spotify_login_required
def top_artists(spotify: spotipy.Spotify):
    db_sess = db_session.create_session()
    params = {'artists': spotify.current_user_top_artists(time_range='short_term'),
              'current_user': db_sess.query(User).get(current_user.id),
              'spotify': spotify,
              }

    return render_template('top_artists.html', **params)

@blueprint.route('/music/test_page')
@login_required
@spotify_login_required
def test_page(spotify: spotipy.Spotify):
    ans = ''
    pprint(spotify.new_releases('RU', limit=1))
    return jsonify(
        {
            'ans': ans
        }
    )

@blueprint.route('/spotify_sign_out')
@spotify_login_required
def sign_out(spotify: spotipy.Spotify):
    try:
        # Remove the CACHE file (.cache-test) so that a new user can authorize.
        os.remove(session_cache_path())
        session.clear()
    except OSError as e:
        print("Error: %s - %s." % (e.filename, e.strerror))
    return redirect('/')

@blueprint.route('/playlists')
@spotify_login_required
def playlists(spotify: spotipy.Spotify):
    return spotify.current_user_playlists()

@blueprint.route('/currently_playing')
@spotify_login_required
def currently_playing(spotify: spotipy.Spotify):
    track = spotify.current_user_playing_track()
    if not track is None:
        return track
    return "No track currently playing."

@blueprint.route('/spotify_current_user')
@spotify_login_required
def get_current_user(spotify: spotipy.Spotify):
    return spotify.current_user()

@blueprint.route('/music/play/<uri>', methods=['POST'])
@spotify_login_required
def play(uri, spotify: spotipy.Spotify):
    spotify.start_playback(context_uri=uri)


import os
import random
from tabulate import tabulate
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from torch import optim
from torch.optim.swa_utils import AveragedModel, SWALR
import copy
from copy import deepcopy
import numpy as np
from sklearn import metrics
from tqdm import tqdm
import matplotlib.pyplot as plt

def enable_dropout(model):
    for m in model.modules():
        if m.__class__.__name__.startswith('Dropout'):
            m.train()

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

class MultiStageModel(nn.Module):
    def __init__(self, num_stages, num_layers, num_f_maps, dim, num_classes, arch):
        super(MultiStageModel, self).__init__()
        self.tower_stage = TowerModel(num_layers, num_f_maps, dim, num_classes, arch)
        self.single_stages = nn.ModuleList([copy.deepcopy(SingleStageModel(num_layers, num_f_maps, num_classes, num_classes, 3, arch))
                                     for s in range(num_stages-1)])

    def forward(self, x, mask):
        middle_out, out = self.tower_stage(x, mask)
        outputs = out.unsqueeze(0)
        for s in self.single_stages:
            middle_out, out = s(F.softmax(out, dim=1) * mask[:, 0:1, :], mask)
            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
        return middle_out, outputs

class TowerModel(nn.Module):
    def __init__(self, num_layers, num_f_maps, dim, num_classes, arch):
        super(TowerModel, self).__init__()
        self.stage1 = SingleStageModel(num_layers, num_f_maps, dim, num_classes, 3, arch)
        self.stage2 = SingleStageModel(num_layers, num_f_maps, dim, num_classes, 5, arch)

    def forward(self, x, mask):
        out1, final_out1 = self.stage1(x, mask)
        out2, final_out2 = self.stage2(x, mask)

        return out1 + out2, final_out1 + final_out2

class SingleStageModel(nn.Module):
    def __init__(self, num_layers, num_f_maps, dim, num_classes, kernel_size, arch):
        super(SingleStageModel, self).__init__()
        self.conv_1x1 = nn.Conv1d(dim, num_f_maps, 1)
        if arch == 'casual':
            self.layers = nn.ModuleList([copy.deepcopy(DilatedResidualCasualLayer(2 ** i, num_f_maps, num_f_maps, kernel_size))
                                        for i in range(num_layers)])
        else:
            self.layers = nn.ModuleList([copy.deepcopy(DilatedResidualLayer(2 ** i, num_f_maps, num_f_maps, kernel_size))
                                        for i in range(num_layers)])
        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)

    def forward(self, x, mask):
        out = self.conv_1x1(x)
        for layer in self.layers:
            out = layer(out, mask)
        final_out = self.conv_out(out) * mask[:, 0:1, :]
        return out, final_out

class DilatedResidualLayer(nn.Module):
    def __init__(self, dilation, in_channels, out_channels, kernel_size):
        super(DilatedResidualLayer, self).__init__()
        padding = int(dilation + dilation * (kernel_size - 3) / 2)
        self.conv_dilated = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation)
        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
        self.dropout = nn.Dropout()

    def forward(self, x, mask):
        out = F.relu(self.conv_dilated(x))
        out = self.conv_1x1(out)
        out = self.dropout(out)
        return (x + out) * mask[:, 0:1, :]

class DilatedResidualCasualLayer(nn.Module):
    def __init__(self, dilation, in_channels, out_channels, kernel_size):
        super(DilatedResidualCasualLayer, self).__init__()
        self.padding = 2 * int(dilation + dilation * (kernel_size -3) / 2)
        self.conv_dilated = nn.Conv1d(in_channels, out_channels, kernel_size, padding=0, dilation=dilation)
        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
        self.dropout = nn.Dropout()

    def forward(self, x, mask):
        out = F.pad(x, [self.padding, 0], 'constant', 0)
        out = F.relu(self.conv_dilated(out))
        out = self.conv_1x1(out)
        out = self.dropout(out)
        return (x + out) * mask[:, 0:1, :]

class Trainer:
    def __init__(self, test_features, test_gt_path, phase2label, device, num_blocks, num_layers, num_f_maps, dim, num_classes, args):
        self.test_features = test_features
        self.test_gt_path = test_gt_path
        self.phase2label = phase2label
        self.device = device
        #self.model = MultiStageModel(num_blocks, num_layers, num_f_maps, dim, num_classes, args.arch)
        self.arch = args.arch
        self.num_blocks = num_blocks
        self.num_layers = num_layers
        self.num_f_maps = num_f_maps
        self.dim = dim
        self.ce = nn.CrossEntropyLoss(ignore_index=-100)
        self.mse = nn.MSELoss(reduction='none')
        self.kl = nn.KLDivLoss(reduction='none')
        self.num_classes = num_classes

        self.pseudo = args.pseudo
        self.num_epochs = args.num_epochs
        self.noisy = args.noisy
        self.uncertainty_warmup_epochs = args.uncertainty_warmup_epochs
        self.max_thres = args.max_thres
        self.smooth = args.smooth
        self.lambda_smooth = args.lambda_smooth
        self.entropy = args.entropy
        self.lambda_entropy = args.lambda_entropy
        self.forward_times = args.forward_times
        self.beta = 0.99

        self.visualization = args.visualization

    def adjust_uncertainty_thres(self, epoch):
        if self.num_classes == 7:  # cholec80
            self.thres = 0.1
        else:
            self.thres = 0.1

    def confidence_loss(self, pred, confidence_mask):
        batch_size = pred.size(0)
        pred = F.log_softmax(pred, dim=1)
        loss = 0
        for b in range(batch_size):
            num_frame = confidence_mask[b].shape[2]
            m_mask = torch.from_numpy(confidence_mask[b]).type(torch.float).to(self.device)
            left = pred[b, :, 1:] - pred[b, :, :-1]
            left = torch.clamp(left[:, :num_frame] * m_mask[0], min=0)
            left = torch.sum(left) / torch.sum(m_mask[0])
            loss += left

            right = (pred[b, :, :-1] - pred[b, :, 1:])
            right = torch.clamp(right[:, :num_frame] * m_mask[1], min=0)
            right = torch.sum(right) / torch.sum(m_mask[1])
            loss += right

        return loss

    def class_balanced_loss(self, logit, target):
        samples_per_cls = []
        for l in range(self.num_classes):
            samples_per_cls.append(int(torch.sum(target==l)))
        effective_num = 1.0 - np.power(self.beta, samples_per_cls)
        weights = (1.0 - self.beta) / np.array(effective_num)
        weights = torch.tensor(weights).float().to(self.device)
        prob = logit.log_softmax(dim=1)
        cb_loss = F.cross_entropy(input=prob, target=target, weight=weights, ignore_index=-100)
        return cb_loss

    def uncertainty_loss(self, logits, thres):
        prob = F.softmax(logits.transpose(2, 1), dim=1)
        entropy = torch.sum(- prob * torch.log(prob), dim=2)
        mask = entropy < thres
        return (entropy * mask).mean()

    def uncertainty_loss_star(self, logits, pseudo_label):
        prob = F.softmax(logits.transpose(2, 1), dim=2)
        entropy = torch.sum(- prob * torch.log(prob), dim=2)
        mask = (pseudo_label != -100)
        return (entropy*entropy * mask).sum() / mask.sum()

    def uncertainty_loss_score(self, logits, scores):
        prob = F.softmax(logits.transpose(2, 1), dim=2)
        entropy = torch.sum(- prob * torch.log(prob), dim=2)
        scores = torch.sigmoid(scores)
        return (entropy * scores).sum() / scores.sum()

    def entropy_regularization(self, logits, mask):
        prob = F.softmax(logits, dim=1)
        entropy = torch.sum(-prob * torch.log(prob), dim=1)
        entropy = entropy * mask
        return torch.sum(entropy) / torch.sum(mask)

    def multi_train(self, save_dir, batch_gen, batch_size, learning_rate):
        # self.train(save_dir, batch_gen, batch_size, learning_rate)
        if self.num_classes == 7:
            times = 3
        else:
            times = 6
        for i in range(times):
            self._multi_train(save_dir, batch_gen, batch_size, learning_rate, i)

    def label_dropout(self, batch_pseudo_target, p):
        pseudo_target = batch_pseudo_target.clone()
        for b in range(pseudo_target.size(0)):
            for f in range(pseudo_target.size(1)):
                if pseudo_target[b, f] != -100:
                    r = random.random()
                    pseudo_target[b, f] = -100 if r<=p else pseudo_target[b, f]
        return pseudo_target

    def mask_boundary(self, batch_target):
        ret_target = deepcopy(batch_target)
        for i in range(batch_target.size(0)):
            boundaries = []
            for j in range(batch_target.size(1)-1):
                if batch_target[i, j] != batch_target[i, j+1]:
                    boundaries.append(j+1)
            for b in boundaries:
                ret_target[i, b-10:b+10] = -100
        return ret_target

    def add_noisy_labels(self, pseudo_target, timestamps, seq_lengths, noisy_length=10, C=0.5):
        from scipy import stats
        norm = stats.norm(0, 1)
        pseudo_target = pseudo_target.clone()
        for b in range(len(pseudo_target)):
            timestamp = timestamps[b]
            seq_length = seq_lengths[b]
            for i in range(len(timestamp)):
                pos = timestamp[i]
                L = 0 if i==0 else timestamp[i-1]+1
                R = seq_length-1 if i==len(timestamp)-1 else timestamp[i+1]-1
                left = min(max(0, pos-L-noisy_length), noisy_length)
                right = min(max(R-pos-noisy_length, 1), noisy_length)
                positions = np.array(list(range(pos-left, pos)) + list(range(pos+1, pos+right)))
                if len(positions) == 0:
                    continue
                probabilites = np.array([norm.pdf(np.abs(p - pos) / noisy_length) for p in positions])
                probabilites = probabilites / np.sum(probabilites)
                choice = np.random.choice(positions, int(len(positions)*C), p=probabilites, replace=False)
                pseudo_target[b][choice.tolist()] = int(pseudo_target[b][timestamp[i]])
        return pseudo_target

    def _multi_train(self, save_dir, batch_gen, batch_size, learning_rate, iter_times):
        self.model = MultiStageModel(self.num_blocks, self.num_layers, self.num_f_maps, self.dim, self.num_classes, self.arch)
        self.model.to(self.device)
        swa_model = AveragedModel(self.model)
        optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.num_epochs, eta_min=0)
        swa_start = 45
        swa_scheduler = SWALR(optimizer, swa_lr=1e-5)
        #scheduler = CosineAnnealingWarmupRestarts(optimizer, first_cycle_steps=self.num_epochs, cycle_mult=1.0, max_lr=learning_rate, min_lr=1e-5, warmup_steps=5, gamma=1.0)
        writer = SummaryWriter(save_dir)
        total_lengths = 0
        labeled_lengths = 0
        correct_lengths = 0
        self.adjust_uncertainty_thres(iter_times)
        for epoch in range(1, self.num_epochs+1):
            #if iter_times == 0:
            #    batch_gen.update_pseudo_labels_noisy()
            self.model.train()
            epoch_loss = 0
            correct = 0
            total = 0
            if self.pseudo == 'confidence' and epoch == self.num_epochs:
                confidence_labels = {}
            while batch_gen.has_next():
                batch_input, batch_target, batch_pseudo_target, mask, batch_confidence = batch_gen.next_batch(batch_size)
                batch_input, batch_target, batch_pseudo_target, mask = batch_input.to(self.device), batch_target.to(self.device),\
                        batch_pseudo_target.to(self.device), mask.to(self.device)
                optimizer.zero_grad()
                middle_pred, predictions = self.model(batch_input, mask)
                timestamp_mask = batch_gen.get_single_random(batch_input.size(-1)).to(self.device)
                vids, seq_lengths, timestamps = batch_gen.get_useful_info()
                # batch_pseudo_target = self.label_dropout(batch_pseudo_target, 0.5)
                if self.noisy:
                    batch_pseudo_target = self.add_noisy_labels(batch_pseudo_target, timestamps, seq_lengths)
                if epoch == self.num_epochs:
                    total_lengths += sum(seq_lengths)
                    labeled_lengths += torch.sum(batch_pseudo_target != -100)
                    for b, seq_length in enumerate(seq_lengths):
                        correct_lengths += torch.sum(batch_pseudo_target[b, :seq_length] == batch_target[b, :seq_length])
                loss = 0
                if self.visualization and self.pseudo == 'confidence' and epoch == self.num_epochs:
                    pseudo_labels = batch_gen.get_boundary(batch_input.size(0), middle_pred)
                    vis_dir = os.path.join(save_dir, 'train_vis')
                    if not os.path.exists(vis_dir):
                        os.makedirs(vis_dir)
                    for i, vi in enumerate(vids):
                        pseudo_label = pseudo_labels[i].squeeze(0)
                        gt = batch_target[i].squeeze(0)
                        length = seq_lengths[i]
                        predicted = torch.max(predictions[-1].data, 1)[1][i].squeeze(0)
                        self.segment_bars(os.path.join(vis_dir, '{}.png'.format(vi)), ['gt', gt[:length]], ['pseudo label', pseudo_label[:length]], ['predict', predicted[:length]])
                if self.pseudo == 'confidence' and epoch == self.num_epochs:
                    pseudo_labels = batch_gen.get_boundary(batch_input.size(0), middle_pred)
                    for i, vi in enumerate(vids):
                        confidence_labels[vi] = pseudo_labels[i].squeeze(0)[:seq_lengths[i]]
                for pi, p in enumerate(predictions):
                    if self.pseudo == 'uncertainty' or self.pseudo == 'self':
                        target = batch_pseudo_target
                    elif self.pseudo == 'uniform':
                        target = batch_gen.get_average(batch_input.size(2)).to(self.device)
                    elif self.pseudo == 'confidence':
                        if epoch > 30:
                            target = batch_gen.get_boundary(batch_input.size(0), middle_pred.detach()).to(self.device)
                        else:
                            target = batch_pseudo_target
                    else:
                        target = batch_target
                        # target = self.mask_boundary(target)
                    ce_loss = self.ce(p.transpose(2, 1).contiguous().view(-1, self.num_classes), target.view(-1))
                    loss += ce_loss
                    #cb_loss = self.class_balanced_loss(p.transpose(2, 1).contiguous().view(-1, self.num_classes), target.view(-1))
                    #loss += cb_loss
                    #print('ce', cs_loss)
                    if self.smooth:
                        smooth_loss = self.lambda_smooth * torch.mean(torch.clamp(
                            self.mse(F.log_softmax(p[:, :, 1:], dim=1), F.log_softmax(p.detach()[:, :, :-1], dim=1)), min=0,
                            max=16) * mask[:, :, 1:])
                        #smooth_loss = self.lambda_smooth * torch.mean(
                        #    self.kl(F.log_softmax(p[:, :, 1:], dim=1), F.log_softmax(p.detach()[:, :, :-1], dim=1)) * mask[:, :, 1:])
                        loss += smooth_loss
                    if self.pseudo == 'confidence' and epoch > 30:
                        confidence_loss = 0.075 * self.confidence_loss(p, batch_confidence)
                        loss += confidence_loss
                    if self.entropy:
                        labeled_mask = target.view(-1) != -100
                        entropy_loss = 0.1 * self.entropy_regularization(p.transpose(2, 1).contiguous().view(-1, self.num_classes), labeled_mask)
                        loss += entropy_loss
                epoch_loss += loss.item()
                loss.backward()
                optimizer.step()

                _, predicted = torch.max(predictions[-1].data, 1)
                correct += ((predicted == batch_target).float()*mask[:, 0, :].squeeze(1)).sum().item()
                total += torch.sum(mask[:, 0, :]).item()
            if self.pseudo == 'confidence' and epoch == self.num_epochs:
                np.save(os.path.join(save_dir, 'pseudo_labels.npy'), confidence_labels, allow_pickle=True)
            if epoch > swa_start:
                swa_model.update_parameters(self.model)
                swa_scheduler.step()
            else:
                scheduler.step()
            batch_gen.reset()
            print('Iter {} Epoch {} training acc :{}'.format(iter_times, epoch, correct / total))

            torch.save(self.model.state_dict(), save_dir + "/epoch-" + str(epoch) + ".model")
            torch.save(optimizer.state_dict(), save_dir + "/epoch-" + str(epoch) + ".opt")
            # self.predict(save_dir, epoch, batch_gen.sample_rate)
        print('Pseudo labeling rate: {:.4f}, pseudo labels accuracy: {:.4f}'.format(labeled_lengths / total_lengths, correct_lengths / labeled_lengths))
        print('Uncertainty threshold: {:.4f}'.format(self.thres))
        self.predict(save_dir, self.num_epochs, batch_gen.sample_rate)
        self.model.eval()
        '''if self.pseudo == 'self':
            with torch.no_grad():
                while batch_gen.has_next():
                    batch_input, _, _, mask = batch_gen.next_batch(batch_size)
                    batch_input, mask = batch_input.to(self.device), mask.to(self.device)
                    _, logits = self.model(batch_input, mask)
                    batch_gen.update_self_labels(logits[-1].detach().transpose(2,1))
                batch_gen.reset()
            return'''
        enable_dropout(self.model)
        with torch.no_grad():
            while batch_gen.has_next():
                batch_input, _, _, mask, _ = batch_gen.next_batch(batch_size)
                batch_input, mask = batch_input.to(self.device), mask.to(self.device)
                mc_probs = []
                for _ in range(self.forward_times):
                    _, logits = self.model(batch_input, mask)
                    probs = F.softmax(logits[-1].squeeze(0).detach().transpose(2, 1), dim=2)
                    mc_probs.append(probs)
                mc_probs = torch.stack(mc_probs, dim=0)
                std = torch.std(mc_probs, dim=0)
                mc_probs = torch.mean(mc_probs, dim=0)
                predictions = torch.max(mc_probs, dim=2)[1]
                uncertainty_scores = std.gather(2, predictions.unsqueeze(-1))
                # uncertainty_scores = torch.sum(-mc_probs * torch.log(mc_probs), dim=2)
                batch_gen.update_pseudo_labels(predictions, uncertainty_scores, thres=self.thres)
                #batch_gen.update_pseudo_labels_all(predictions, uncertainty_scores, thres=self.thres)
                if False and self.visualization:
                    data_dir = os.path.join(save_dir, 'train_epoch{}_score'.format(epoch))
                    if not os.path.exists(data_dir):
                        os.makedirs(data_dir)
                    vids, seq_lengths, timestamp_pos = batch_gen.get_useful_info()
                    for b, score in enumerate(uncertainty_scores):
                        np.save(os.path.join(data_dir, '{}.npy'.format(vids[b])), {'timestamp': timestamp_pos[b], 'score': score[:seq_lengths[b]].cpu().numpy()})
            batch_gen.reset()

            if self.visualization:
                self.visualize(batch_gen, save_dir, iter_times, batch_size)
        batch_gen.save_pseudo_labels(save_dir, iter_times)

    def train(self, save_dir, batch_gen, batch_size, learning_rate):
        self.model.to(self.device)
        optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.num_epochs, eta_min=1e-5)
        #scheduler = CosineAnnealingWarmupRestarts(optimizer, first_cycle_steps=self.num_epochs, cycle_mult=1.0, max_lr=learning_rate, min_lr=1e-5, warmup_steps=5, gamma=1.0)
        writer = SummaryWriter(save_dir)
        for epoch in range(1, self.num_epochs+1):
            self.adjust_uncertainty_thres(epoch)
            lr = scheduler.get_lr()[0]
            thres = self.thres
            print('='*20 + 'Epoch {}'.format(epoch) + '='*20)
            print('learning rate = {:.4f}, uncertainty thresold = {:.4f}'.format(lr, thres))
            self.model.train()
            epoch_loss = 0
            correct = 0
            total = 0
            while batch_gen.has_next():
                batch_input, batch_target, batch_pseudo_target, mask = batch_gen.next_batch(batch_size)
                batch_input, batch_target, batch_pseudo_target, mask = batch_input.to(self.device), batch_target.to(self.device),\
                        batch_pseudo_target.to(self.device), mask.to(self.device)
                optimizer.zero_grad()
                _, predictions = self.model(batch_input, mask)
                timestamp_mask = batch_gen.get_single_random(batch_input.size(-1)).to(self.device)
                _, seq_lengths, _ = batch_gen.get_useful_info()
                total_lengths = sum(seq_lengths)
                labeled_lengths = torch.sum(batch_pseudo_target != -100)
                loss = 0
                for pi, p in enumerate(predictions):
                    if self.pseudo == 'full':
                        target = batch_target
                    else:
                        target = batch_pseudo_target
                    #ce_loss = self.ce(p.transpose(2, 1).contiguous().view(-1, self.num_classes), target.view(-1))
                    #loss += ce_loss
                    cb_loss = self.class_balanced_loss(p.transpose(2, 1).contiguous().view(-1, self.num_classes), target.view(-1))
                    loss += cb_loss
                    #print('ce', cs_loss)
                    if self.smooth:
                        smooth_loss = self.lambda_smooth * torch.mean(torch.clamp(
                            self.mse(F.log_softmax(p[:, :, 1:], dim=1), F.log_softmax(p.detach()[:, :, :-1], dim=1)), min=0,
                            max=16) * mask[:, :, 1:])
                        loss += smooth_loss

                    if self.entropy:
                        entropy_loss = 0.1 * self.entropy_regularization(p.transpose(2, 1).contiguous().view(-1, self.num_classes))
                        loss += entropy_loss
                #if epoch >= self.start_epoch:
                    #print('ce loss: {:.4f}, smooth_loss: {:.4f}, uncertainty_loss: {:.4f}'.format(ce_loss, smooth_loss, uncertainty_loss))

                epoch_loss += loss.item()
                loss.backward()
                optimizer.step()

                _, predicted = torch.max(predictions[-1].data, 1)
                correct += ((predicted == batch_target).float()*mask[:, 0, :].squeeze(1)).sum().item()
                total += torch.sum(mask[:, 0, :]).item()
            scheduler.step()
            batch_gen.reset()

            if self.pseudo == 'uncertainty' and epoch % 20 == 0 and epoch != self.num_epochs:
                self.model.eval()
                enable_dropout(self.model)
                with torch.no_grad():
                    while batch_gen.has_next():
                        batch_input, _, _, mask = batch_gen.next_batch(batch_size)
                        batch_input, mask = batch_input.to(self.device), mask.to(self.device)
                        mc_probs = []
                        for _ in range(self.forward_times):
                            _, logits = self.model(batch_input, mask)
                            probs = F.softmax(logits[-1].squeeze(0).detach().transpose(2, 1), dim=2)
                            mc_probs.append(probs)
                        mc_probs = torch.stack(mc_probs, dim=0)
                        std = torch.std(mc_probs, dim=0)
                        mc_probs = torch.mean(mc_probs, dim=0)
                        predictions = torch.max(mc_probs, dim=2)[1]
                        # uncertainty_scores = std.gather(2, predictions.unsqueeze(-1))
                        uncertainty_scores = torch.sum(-mc_probs * torch.log(mc_probs), dim=2)
                        batch_gen.update_pseudo_labels(predictions, uncertainty_scores, thres=self.thres)
                        #batch_gen.update_pseudo_labels_all(predictions, uncertainty_scores, thres=self.thres)
                        if self.visualization:
                            data_dir = os.path.join(save_dir, 'train_epoch{}_score'.format(epoch))
                            if not os.path.exists(data_dir):
                                os.makedirs(data_dir)
                            vids, seq_lengths, timestamp_pos = batch_gen.get_useful_info()
                            for b, score in enumerate(uncertainty_scores):
                                np.save(os.path.join(data_dir, '{}.npy'.format(vids[b])), {'timestamp': timestamp_pos[b], 'score': score[:seq_lengths[b]].cpu().numpy()})
                    batch_gen.reset()

                    if self.visualization:
                        self.visualize(batch_gen, save_dir, epoch, batch_size)

            torch.save(self.model.state_dict(), save_dir + "/epoch-" + str(epoch) + ".model")
            torch.save(optimizer.state_dict(), save_dir + "/epoch-" + str(epoch) + ".opt")
            writer.add_scalar('trainLoss', epoch_loss / len(batch_gen.list_of_samples), epoch)
            writer.add_scalar('trainAcc', float(correct)/total, epoch)
            print("Training loss = %f,   acc = %f" % (epoch_loss / len(batch_gen.list_of_samples),
                                                               float(correct)/total))

            self.predict(save_dir, epoch, batch_gen.sample_rate)

    def visualize(self, batch_gen, save_dir, epoch, batch_size):
        vis_dir = os.path.join(save_dir, 'train_epoch{}_seg'.format(epoch))
        if not os.path.exists(vis_dir):
            os.makedirs(vis_dir)
        while batch_gen.has_next():
            batch_input, batch_target, batch_pseudo_target, mask, _ = batch_gen.next_batch(batch_size)
            batch_input, mask = batch_input.to(self.device), mask.to(self.device)
            _, logits = self.model(batch_input, mask)
            vids, seq_lengths, timestamp_pos = batch_gen.get_useful_info()
            for vi in range(len(batch_input)):
                seq_length = seq_lengths[vi]
                pseudo_target = batch_pseudo_target[vi][:seq_length]
                target = batch_target[vi][:seq_length]
                predicted = torch.max(logits[-1].data, dim=1)[1][vi].squeeze()[:seq_length]
                pseudo_target[pseudo_target==-100] = 8  # 11
                #target[timestamp_pos[vi]] = 9  # 10
                assert len(pseudo_target) == len(target) and len(target) == len(predicted)
                self.segment_bars(os.path.join(vis_dir, '{}.png'.format(vids[vi])), ['gt', target], ['pseudo label', pseudo_target], ['predict', predicted])
        batch_gen.reset()

    def predict(self, model_dir, epoch, sample_rate):
        label2phase_dicts = {
            'cholec80':{
                0: 'Preparation',
                1: 'CalotTriangleDissection',
                2: 'ClippingCutting',
                3: 'GallbladderDissection',
                4: 'GallbladderPackaging',
                5: 'CleaningCoagulation',
                6: 'GallbladderRetraction'},
            'm2cai16':{
                0: 'TrocarPlacement',
                1: 'Preparation',
                2: 'CalotTriangleDissection',
                3: 'ClippingCutting',
                4: 'GallbladderDissection',
                5: 'GallbladderPackaging',
                6: 'CleaningCoagulation',
                7: 'GallbladderRetraction'}
        }
        res_dir = os.path.join(model_dir, 'predict')
        if not os.path.exists(res_dir):
            os.makedirs(res_dir)
        self.model.eval()
        with torch.no_grad():
            self.model.to(self.device)
            self.model.load_state_dict(torch.load(model_dir + "/epoch-" + str(epoch) + ".model"))
            videos = [os.path.join(self.test_features, x) for x in sorted(os.listdir(self.test_features))]
            annotations = [os.path.join(self.test_gt_path, x) for x in sorted(os.listdir(self.test_gt_path))]
            vis_dir = os.path.join(model_dir, 'epoch{}_test_vis'.format(epoch))
            if not os.path.exists(vis_dir):
                os.makedirs(vis_dir)
            vid_pre = 41 if  len(videos) == 40 else 28
            label2phase = label2phase_dicts['cholec80' if len(videos)==40 else 'm2cai16']
            #all_pred_phase = []
            #all_label_phase = []
            #correct_phase = 0
            #total_phase = 0

            Acc = AverageMeter()
            Pre = AverageMeter()
            Rec = AverageMeter()
            Jac = AverageMeter()

            for vid, (video, anno) in enumerate(list(zip(videos, annotations))):
                # print(vid)
                features = np.load(video).transpose()
                features = features[:, ::sample_rate]
                with open(anno, 'r') as f:
                    content = f.read().split('\n')
                    if content[-1] == '':
                        content = content[:-1]
                labels = np.zeros(len(content))
                for i in range(len(content)):
                    labels[i] = self.phase2label[content[i].strip().split()[1]]
                labels = torch.Tensor(labels[::sample_rate]).long().to(self.device)

                input_x = torch.tensor(features, dtype=torch.float)
                input_x.unsqueeze_(0)
                input_x = input_x.to(self.device)
                _, predictions = self.model(input_x, torch.ones(input_x.size(), device=self.device))
                _, predicted = torch.max(predictions[-1].data, 1)
                predicted = predicted.squeeze()

                # self.segment_bars(os.path.join(vis_dir, '{}.png'.format(video.split('/')[-1].split('.')[0])), ['gt', labels], ['predict', predicted])

                correct_phase = torch.sum(predicted == labels)
                total_phase = len(predicted)

                pred_phase = []
                label_phase = []
                for i in range(len(predicted)):
                    pred_phase.append(int(predicted.data.cpu()[i]))
                for i in range(len(labels)):
                    label_phase.append(int(labels[i]))
                with open(os.path.join(res_dir, f'video{vid+vid_pre}_pred.txt'), 'w') as f:
                    f.write('Frame\tPhase\n')
                    for fid, each_pred in enumerate(predicted):
                        f.write('{}\t{}\n'.format(fid, label2phase[int(each_pred)]))

                accuracy = correct_phase / total_phase
                precision = metrics.precision_score(label_phase, pred_phase, average='macro')
                recall = metrics.recall_score(label_phase, pred_phase, average='macro')
                jaccard = metrics.jaccard_score(label_phase, pred_phase, average='macro')
                #F1 = metrics.f1_score(label_phase, pred_phase, average='macro')
                Acc.update(accuracy)
                Pre.update(precision)
                Rec.update(recall)
                Jac.update(jaccard)

        #accuracy = correct_phase / total_phase
        #precision = metrics.precision_score(all_label_phase, all_pred_phase, average='macro')
        #recall = metrics.recall_score(all_label_phase, all_pred_phase, average='macro')
        #jaccard = metrics.jaccard_score(all_label_phase, all_pred_phase, average='macro')
        #F1 = metrics.f1_score(all_label_phase, all_pred_phase, average='macro')
        print('Evaluating from {} at epoch {}'.format(model_dir, epoch))
        print(tabulate([['{:.2f}'.format(Acc.avg*100), '{:.2f}'.format(Jac.avg*100), '{:.2f}'.format(Pre.avg*100), '{:.2f}'.format(Rec.avg*100)]],
                    headers=['Accuracy', 'Jaccard', 'Precision', 'Recall'], tablefmt='orgtbl'))
        #print('F1 score: {:.4f}'.format(F1))

    def segment_bars_(self, save_path, *labels):
        color_map = [
                (255, 0, 0),    # red
                (255, 165, 0),  # orange
                (255, 255, 0),  # yellow
                (0, 255, 0),    # green
                (0, 127, 255),  # cyan
                (0, 0, 255),    # blue
                (139, 0, 255),  # purple
                (87, 105, 60),  # dark green
                (0, 0, 0),      # black -> background
                (255, 255, 255),# white -> timestamp
                ]
        titles, labels = zip(*labels)
        labels = [label.detach().cpu().numpy() for label in labels]
        nrows = len(labels)
        cmap = plt.cm.Paired
        barprops = dict(aspect='auto', interpolation='none', cmap=cmap)
        figh = 0.15 + 0.15 + (nrows + (nrows-1)*0.15)*0.5
        fig, axs = plt.subplots(nrows=nrows, figsize=(7, figh))
        fig.subplots_adjust(top=1-.35/figh, bottom=.15/figh, left=0.2, right=0.99)
        for ax, title, label in zip(axs, titles, labels):
            label = np.vstack([label, label])
            #label = (label+0.5) / len(cmap.colors)
            ax.text(-.01, .5, title, va='center', ha='right', fontsize=10, transform=ax.transAxes)
            ax.imshow(label, **barprops)
            ax.set_axis_off()
        fig.savefig(save_path)
        plt.close(fig)

    def segment_bars(self, save_path, *labels):
        def scale_lightness(rgb, scale_l):
            import colorsys
            h, l, s = colorsys.rgb_to_hls(*rgb)
            return colorsys.hls_to_rgb(h, min(1, l * scale_l), s=s)
        color_map = [
                (255, 0, 0),    # red
                (255, 165, 0),  # orange
                (255, 255, 0),  # yellow
                (0, 255, 0),    # green
                (0, 127, 255),  # cyan
                (0, 0, 255),    # blue
                (139, 0, 255),  # purple
                (87, 105, 60),  # dark green
                (0, 0, 0),      # black -> background
                (255, 255, 255),# white -> timestamp
                ]
        color_map = list(plt.cm.Paired.colors[:8])
        color_map.append((1, 1, 1))
        titles, labels = zip(*labels)
        labels = [label.detach().cpu().numpy().astype(np.int32) for label in labels]
        nrows = len(labels)
        figh = 0.15 + 0.15 + (nrows + (nrows-1)*0.15)*0.5
        fig, axs = plt.subplots(nrows=nrows, figsize=(7, figh))
        fig.subplots_adjust(top=1-.35/figh, bottom=.15/figh, left=0.2, right=0.99)
        for ax, title, label in zip(axs, titles, labels):
            #width = 1 / len(label)
            #x = np.linspace(width/2, 1-width/2, len(label))
            width = 1
            x = np.linspace(0.5, len(label)-0.5, len(label))
            ax.text(-.01, .5, title, va='center', ha='right', fontsize=10, transform=ax.transAxes)
            timestamp = []
            for i, (x_, l_) in enumerate(zip(x, label)):
                if l_ == 9:
                    color = color_map[label[i-1] if i>0 else label[i+1]]
                    timestamp.append((x_, color))
                    # ax.bar([x_], [1], width=width, color=scale_lightness(color, 1.7))
                else:
                    ax.bar([x_], [1], width=width, color=color_map[l_], alpha=0.6)
            for x_, color in timestamp:
                ax.bar([x_], [1], width=width*10, color=scale_lightness(color, 0.8))
            ax.set_axis_off()
        fig.savefig(save_path)
        plt.close(fig)


import numpy as np
import matplotlib.pyplot as plt
from numpy import genfromtxt

# Import data

data = genfromtxt('/Users/mikelevine/PycharmProjects/ENSO_VS_HURRICANES/ENSO_VS_NAMED_HURRICANES.csv',
                  delimiter=',')

# Define arrays and remove nan values

hurr = np.zeros(len(data))

for i in range(len(data)):
    hurr[i] = data[i][12]

hurr = hurr[~np.isnan(hurr)]

ensomonths = np.delete(data,[12,13,14,15,16,17,18,19],1)

for i in range(4):
    ensomonths = np.delete(ensomonths,len(ensomonths)-1,0)

# Hurricane Season on Average, runs from June to November (eliminate other months and average enso signal)

ensohseason = np.delete(ensomonths, 11, 1)
for i in range(5):
    ensohseason = np.delete(ensohseason, 0, 1)

# Average hurricane season ENSO signal for each year

ensohseasonmean = np.mean(ensohseason, 1)

fig, ax = plt.subplots()
lns1=ax.plot(np.arange(1950,2015,1), ensohseasonmean, 'r-', label="Mean ENSO Signal (June-November)")
lns2=ax.plot(np.arange(1950,2015,1), hurr, 'b-', label="Number of Hurricanes")
ax.set_xlabel("Year")
ax.set_ylabel("Mean ENSO Signal / Number of Hurricanes")
lns = lns1+lns2
labs = [l.get_label() for l in lns]
legend1 = ax.legend(lns,labs,loc=0, prop={'size':10})
plt.title("Annual Mean ENSO Signal vs Number of Named Hurricanes")
plt.show()

# Define and plot derivatives

densohseasonmean = np.gradient(ensohseasonmean)
dhurr = np.gradient(hurr)

fig, ax = plt.subplots()
lns1=ax.plot(np.arange(1950,2015,1), densohseasonmean, 'r-', label="Change in Mean ENSO Signal")
lns2=ax.plot(np.arange(1950,2015,1), dhurr, 'b-', label="Change in Number of Hurricanes")
ax.set_xlabel("Year")
ax.set_ylabel("Change in Mean ENSO Signal / Number of Hurricanes")
lns = lns1+lns2
labs = [l.get_label() for l in lns]
legend1 = ax.legend(lns,labs,loc=0, prop={'size':10})
plt.title("Change in Annual Mean ENSO Signal vs Number of Named Hurricanes")
plt.show()


# -*- coding: utf-8 -*-
import logging
import werkzeug
import openerp
from openerp import SUPERUSER_ID
from openerp import http
from openerp import tools
from openerp.http import request
from openerp.tools.translate import _
from openerp.addons.website.models.website import slug

PPG = 9 # Products Per Page
PPR = 3  # Products Per Row

_logger = logging.getLogger(__name__)

class table_compute(object):
    def __init__(self):
        self.table = {}

    def _check_place(self, posx, posy, sizex, sizey):
        res = True
        for y in range(sizey):
            for x in range(sizex):
                if posx+x>=PPR:
                    res = False
                    break
                row = self.table.setdefault(posy+y, {})
                if row.setdefault(posx+x) is not None:
                    res = False
                    break
            for x in range(PPR):
                self.table[posy+y].setdefault(x, None)
        return res

    def process(self, products, ppg=PPG):
        # Compute products positions on the grid
        minpos = 0
        index = 0
        maxy = 0
        for p in products:
            x = min(max(p.website_size_x, 1), PPR)
            y = min(max(p.website_size_y, 1), PPR)
            if index>=ppg:
                x = y = 1

            pos = minpos
            while not self._check_place(pos%PPR, pos/PPR, x, y):
                pos += 1
            # if 21st products (index 20) and the last line is full (PPR products in it), break
            # (pos + 1.0) / PPR is the line where the product would be inserted
            # maxy is the number of existing lines
            # + 1.0 is because pos begins at 0, thus pos 20 is actually the 21st block
            # and to force python to not round the division operation
            if index >= ppg and ((pos + 1.0) / PPR) > maxy:
                break

            if x==1 and y==1:   # simple heuristic for CPU optimization
                minpos = pos/PPR

            for y2 in range(y):
                for x2 in range(x):
                    self.table[(pos/PPR)+y2][(pos%PPR)+x2] = False
            self.table[pos/PPR][pos%PPR] = {
                'product': p, 'x':x, 'y': y,
                'class': " ".join(map(lambda x: x.html_class or '', p.website_style_ids))
            }
            if index<=ppg:
                maxy=max(maxy,y+(pos/PPR))
            index += 1

        # Format table according to HTML needs
        rows = self.table.items()
        rows.sort()
        rows = map(lambda x: x[1], rows)
        for col in range(len(rows)):
            cols = rows[col].items()
            cols.sort()
            x += len(cols)
            rows[col] = [c for c in map(lambda x: x[1], cols) if c != False]

        return rows

class QueryURL(object):
    def __init__(self, path='', **args):
        self.path = path
        self.args = args

    def __call__(self, path=None, **kw):
        if not path:
            path = self.path
        for k,v in self.args.items():
            kw.setdefault(k,v)
        l = []
        for k,v in kw.items():
            if v:
                if isinstance(v, list) or isinstance(v, set):
                    l.append(werkzeug.url_encode([(k,i) for i in v]))
                else:
                    l.append(werkzeug.url_encode([(k,v)]))
        if l:
            path += '?' + '&'.join(l)
        return path

class Loewie_Download_Ctl(http.Controller):
    """
    @http.route('/web/login', type='http', auth="none")
    def web_login(self, redirect=None, **kw):
        #ensure_db()
        request.params['login_success'] = False
        if request.httprequest.method == 'GET' and redirect and request.session.uid:
            return http.redirect_with_hash(redirect)

        if not request.uid:
            request.uid = openerp.SUPERUSER_ID

        values = request.params.copy()
        try:
            values['databases'] = http.db_list()
        except openerp.exceptions.AccessDenied:
            values['databases'] = None

        values['error'] = _("Please Login to See our download resources.")
        return request.render('web.login', values) """

    @http.route([
        '/downloads',
        '/downloads/<model("res.users"):user>',
        '/downloads/page/<int:page>',
        '/downloads/category/<model("download.category"):category>',
        '/downloads/category/<model("download.category"):category>/page/<int:page>'
    ], type='http', auth="user", website=True)
    def downloads(self, page=0, category=None, search='', ppg=False, **post):

        cr, uid, context, pool = request.cr, request.uid, request.context, request.registry
        user_obj = pool['res.users'].browse(cr,uid,[uid],context=context)
        #if not user_obj.active :
        #    request.params['error'] = _("Please Login to See our download resources.")
        #    return request.redirect('/web/login')

        category_obj = pool['download.category']
        category_ids = category_obj.search(cr, uid, [('parent_id', '=', False)], context=context)
        categories = category_obj.browse(cr, uid, category_ids, context=context).sorted(key=lambda x:x.sequence)
        download_obj = pool.get('loewie.download')
        if category and category.name == 'All Downloads' :
            category=None

        keep = QueryURL('/downloads', category=category and int(category), search=search)

        url = "/downloads"
        category_ids = []
        if category:
            category_ids.append(category.id)
            url = "/downloads/category/%s" % slug(category)
            for categ in category.child_id:
                category_ids.append(categ.id)

        domain = []
        if category_ids : domain = [('category_id','in',category_ids or [1])]

        if search :
            for srch in search.split(" "):
                domain += [('name', 'ilike', srch)]

        download_count = download_obj.search_count(cr, uid, domain, context=context)
        pager = request.website.pager(url=url, total=download_count, page=page, step=9, scope=7, url_args=post)
        download_ids = download_obj.search(cr, uid, domain, limit=9, offset=pager['offset'], context=context) # order=self._get_search_order(post),
        #download_ids = download_obj.search(cr, uid, domain, context=context)
        downloads = download_obj.browse(cr, uid, download_ids, context=context)

        parent_category_ids = []
        if category:
            parent_category_ids = [category.id]
            current_category = category
            while current_category.parent_id:
                parent_category_ids.append(current_category.parent_id.id)
                current_category = current_category.parent_id

        values = {
            'search': search,
            'categories': categories,
            'category': category,
            'parent_category_ids':parent_category_ids,
            'keep': keep,
            'pager': pager,
            'downloads': downloads,
            #'bins': table_compute().process(downloads, ppg),
            'rows': 3,
        }

        return request.website.render("nomitang_download.downloads_tmpl", values)

    @http.route(['/downloads/download/<model("loewie.download"):download>', '/downloads/download/<model("res.users"):user>', ], type='http', auth="user", website=True)
    def download(self, download, category='', search='', **kwargs):
        cr, uid, context, pool = request.cr, request.uid, request.context, request.registry
        user_obj = pool['res.users'].browse(cr,uid,[uid],context=context)
        #if not user_obj.active :
        #    return request.redirect('/website/login')

        #domain = []
        if search : request.redirect(request.httprequest.referrer or '/downloads')
            #for srch in search.split(" "):
            #    domain += [('name', 'ilike', srch)]

        category_obj = pool['download.category']
        download_obj = pool['loewie.download']

        context.update(active_id=download.id)

        if category:
            category = category_obj.browse(cr, uid, int(category), context=context)
            category = category if category.exists() else False

        category_ids = category_obj.search(cr, uid, [('parent_id', '=', False)], context=context)
        categories = category_obj.browse(cr, uid, category_ids, context=context).sorted(key=lambda x:x.sequence)
        keep = QueryURL('/downloads', category=category and category.id, search=search)

        parent_category_ids = []
        if category:
            parent_category_ids = [category.id]
            current_category = category
            while current_category.parent_id:
                parent_category_ids.append(current_category.parent_id.id)
                current_category = current_category.parent_id

        values = {
            'search': search,
            'categories': categories,
            'category': category,
            'parent_category_ids':parent_category_ids,
            'main_object': download,
            'keep': keep,
        }
        return request.website.render("nomitang_download.download_tmpl", values)


from turtle import Turtle, Screen
from paddle import Paddle
from ball import Ball
from scoreboard import Scoreboard
import time

screen = Screen()
screen.bgcolor('Black')
screen.setup(width=800, height=600)
screen.title("Pong")
screen.tracer(0)

l_paddle = Paddle((-350, 0))
r_paddle = Paddle((350, 0))
ball = Ball()
score = Scoreboard()

screen.listen()
screen.onkey(l_paddle.go_up, "w")
screen.onkey(l_paddle.go_down, "s")
screen.onkey(r_paddle.go_up, "Up")
screen.onkey(r_paddle.go_down, "Down")

game_is_on = True
sleep = 0.05
while game_is_on:
    time.sleep(sleep)
    screen.update()
    score.show()
    ball.move()
    if score.l_score == 5:
        score.goto(0, 210)
        score.write("Left player wins!!", align="Center", font=("Helvetica", 30, "normal"))
        game_is_on = False

    if score.r_score == 5:
        score.goto(0, 210)
        score.write("Right player wins!!", align="Center", font=("Helvetica", 30, "normal"))
        game_is_on = False

    # detect collision with the wall
    if abs(ball.ycor()) >= 290:
        ball.bounce_y()

    # detect collision with a paddle
    if ball.distance(r_paddle) < 50 and ball.xcor() > 320 or ball.distance(l_paddle) < 50 and ball.xcor() < -320:
        ball.bounce_x()
        sleep *= 0.9

    # detect out of bounce
    if ball.xcor() > 380:
        score.l_score += 1
        ball.reset_position()

    if ball.xcor() < -380:
        score.r_score += 1
        ball.reset_position()

    ball.debounce()

    # debug
    print(ball.ycor(), ball.xcor(), ball.y_move, ball.metric, sleep)

screen.exitonclick()