import os
from pathlib import Path

import win32com.client
import openpyxl
from tkinter import *
from tkinter import filedialog

def xls_to_xlsx(folder_path, file_name):
    """
    excel  .xls 后缀 改成 .xlsx 后缀
    folder_path 文件夹路径
    file_name 文件名字 带后缀 比如 aa.xls
    """
    folder_path = folder_path.replace('/', '\\')
    file_name = file_name.replace('/', '\\')
    name, suffix = file_name.split('.')
    excel_file_path = os.path.join(folder_path, file_name)

    excel = win32com.client.gencache.EnsureDispatch('Excel.Application')  # 要看MIME手册
    wb = excel.Workbooks.Open(excel_file_path)
    suffix = f".{suffix}x"
    new_file_name = f"{name}{suffix}"
    new_excel_file_path = os.sep.join([folder_path, new_file_name])
    # test
    print("new_excel_file_path: " + new_excel_file_path)
    wb.SaveAs(new_excel_file_path, FileFormat=51)
    wb.Close()
    excel.Application.Quit()
    return new_excel_file_path

def generate_one_warrant(excel_path, confirm_sheet, index, para_list):
    warrant_book = openpyxl.load_workbook(excel_path, data_only=True)
    warrant_sheet1 = warrant_book["1"]  # 股权证首页
    warrant_sheet2 = warrant_book["2"]  # 股东基本信息登记
    warrant_sheet3 = warrant_book["3"]  # 股权登记
    # warrant_sheet2['E25'] = confirm_sheet['A10'].value

    confirm_row_index = 10
    warrant_sheet2_row_index = 25
    warrant_sheet3_row_index = 24

    sheet_id = confirm_sheet['K2'].value

    count = 0
    while confirm_sheet.cell(confirm_row_index, 1).value is not None \
            and not (('家庭' in confirm_sheet.cell(confirm_row_index, 1).value)
                     and (confirm_sheet.cell(confirm_row_index, 3).value is None)):
        member_name = confirm_sheet.cell(confirm_row_index, 1).value
        id_number = confirm_sheet.cell(confirm_row_index, 5).value
        gender = ''
        if id_number is None or id_number == '':
            print("警告：编号为{}的表中身份证号缺失".format(sheet_id))
        else:
            id_number = id_number.strip()
            if len(id_number) != 18:
                print("警告：编号为{}的表中身份证号位数不为18".format(sheet_id))
            if len(id_number) < 2:
                print("警告：编号为{}的表中身份证号位数小于2".format(sheet_id))
            else:
                if id_number[-2].isdigit():
                    if int(id_number[-2]) % 2 == 0:
                        gender = '女'
                    else:
                        gender = '男'

        warrant_sheet2.cell(warrant_sheet2_row_index, 5).value = member_name  # 姓名
        warrant_sheet2.cell(warrant_sheet2_row_index, 6).value = gender  # 性别
        warrant_sheet2.cell(warrant_sheet2_row_index, 7).value = id_number  # 身份证号码
        warrant_sheet2.cell(warrant_sheet2_row_index, 8).value = confirm_sheet.cell(confirm_row_index, 3).value  # 与户主关系

        warrant_sheet3.cell(warrant_sheet3_row_index, 19).value = member_name
        warrant_sheet3.cell(warrant_sheet3_row_index, 20).value = 10
        warrant_sheet3.cell(warrant_sheet3_row_index, 23).value = '成员股'

        confirm_row_index = confirm_row_index + 1
        warrant_sheet2_row_index = warrant_sheet2_row_index + 1
        warrant_sheet3_row_index = warrant_sheet3_row_index + 1
        count = count + 1
        # '家庭' not in confirm_sheet.cell(row_index, 1).value) and (confirm_sheet.cell(row_index, 3).value is not None

        if count > 14:
            print('警告：编号为{}的股权证中人数超过14，需手动调整！'.format(sheet_id))

    warrant_sheet3['T39'] = count * 10

    last4 = str(index + 1).zfill(4)

    warrant_sheet1['X22'] = para_list[0]
    warrant_sheet1['V31'] = para_list[1]
    warrant_sheet1['W54'] = para_list[2]
    warrant_sheet1['W81'] = para_list[3] + last4
    warrant_sheet1['V89'] = para_list[4]
    warrant_sheet1['X89'] = para_list[5]
    warrant_sheet1['AA89'] = para_list[6]

    warrant_sheet3['T42'] = para_list[7]

    filename_save = last4 + confirm_sheet['A10'].value + '股权证.xlsx'
    (Path("warrant") / para_list[1]).mkdir(parents=True, exist_ok=True)

    folder_path2 = Path("warrant") / para_list[1]
    whole_save = folder_path2 / filename_save
    warrant_book.save(whole_save)
    # print(count)
    warrant_book.close()

def generate_share_warrants(warrant_path, confirm_path, para_list):
    folder_path, file_name = os.path.split(warrant_path)
    # folder_path = r"C:\Users\sc\Desktop\test_openpyxl"
    # file_name = 'sample.xlsx'
    old_suffix = file_name.split('.')[-1]
    if old_suffix == 'xls':
        print("transform .xls to .xlsx")
        excel_path = xls_to_xlsx(folder_path, file_name)
    elif old_suffix == 'xlsx':
        print("no need to transform file type")
        excel_path = os.sep.join([folder_path, file_name])
        # print(excel_path)
    else:
        print("wrong file type: " + old_suffix)
        return

    folder_path2, file_name2 = os.path.split(confirm_path)
    # folder_path = r"C:\Users\sc\Desktop\test_openpyxl"
    # file_name = 'sample.xlsx'
    old_suffix2 = file_name2.split('.')[-1]
    if old_suffix2 == 'xls':
        print("transform .xls to .xlsx")
        excel_path2 = xls_to_xlsx(folder_path2, file_name2)
    elif old_suffix2 == 'xlsx':
        print("no need to transform file type")
        excel_path2 = os.sep.join([folder_path2, file_name2])
        # print(excel_path)
    else:
        print("wrong file type: " + old_suffix2)
        return

    confirm_book = openpyxl.load_workbook(excel_path2, data_only=True)
    confirm_sheet_names = confirm_book.sheetnames

    for i in range(len(confirm_sheet_names)):
        generate_one_warrant(excel_path, confirm_book[confirm_sheet_names[i]], i, para_list)

    if old_suffix == 'xls':
        os.remove(excel_path)

    confirm_book.close()

if __name__ == '__main__':
    path_input = ''

    top = Tk()
    frame_file = Frame(top)

    # 变量path
    path = StringVar()
    # 输入框，标记，按键
    Label(top, text="确认表路径:").pack()
    # 输入框绑定变量path
    Entry(top, textvariable=path, width=30).pack()

    def selectPath():
        global path_input
        path_input = filedialog.askopenfilename(title='请选择确认登记表', filetypes=[('Excel', '.xls .xlsx')])
        path.set(path_input)
        print(path_input)

    Button(top, text="路径选择", command=selectPath).pack()

    # 1
    frame_credit_code = Frame(top)
    label_credit_code = Label(frame_credit_code, text="统一社会信用代码")
    label_credit_code.pack(side=LEFT)
    entry_credit_code = Entry(frame_credit_code)
    entry_credit_code.pack(side=RIGHT)
    frame_credit_code.pack()

    # 2
    frame_organization = Frame(top)
    label_organization = Label(frame_organization, text="组织名称")
    label_organization.pack(side=LEFT)
    entry_organization = Entry(frame_organization)
    entry_organization.pack(side=RIGHT)
    frame_organization.pack()

    # 3
    frame_representative = Frame(top)
    label_representative = Label(frame_representative, text="法定代表人")
    label_representative.pack(side=LEFT)
    entry_representative = Entry(frame_representative)
    entry_representative.pack(side=RIGHT)
    frame_representative.pack()

    # 4
    frame_number = Frame(top)
    label_number = Label(frame_number, text="股权证编号（除了后四位）")
    label_number.pack(side=LEFT)
    entry_number = Entry(frame_number)
    entry_number.pack(side=RIGHT)
    frame_number.pack()

    # 5, 6, 7
    frame_issue_date = Frame(top)
    Label(frame_issue_date, text="发证日期").pack(side=LEFT)
    entry_year = Entry(frame_issue_date, width=5)
    entry_year.pack(side=LEFT)
    Label(frame_issue_date, text="年").pack(side=LEFT)
    entry_month = Entry(frame_issue_date, width=5)
    entry_month.pack(side=LEFT)
    Label(frame_issue_date, text="月").pack(side=LEFT)
    entry_day = Entry(frame_issue_date, width=5)
    entry_day.pack(side=LEFT)
    Label(frame_issue_date, text="日").pack(side=LEFT)

    frame_issue_date.pack()

    # 8
    frame_register_date = Frame(top)
    Label(frame_register_date, text="登记日期").pack(side=LEFT)
    entry_register_date = Entry(frame_register_date)
    entry_register_date.pack(side=LEFT)
    frame_register_date.pack()

    def btn_generator_CallBack():
        para_list = [entry_credit_code.get(),
                     entry_organization.get(),
                     entry_representative.get(),
                     entry_number.get(),
                     entry_year.get(),
                     entry_month.get(),
                     entry_day.get(),
                     entry_register_date.get()]

        # messagebox.showinfo("Hello Python", para_list)

        print(os.getcwd())
        sample_path = os.path.join(os.getcwd(), 'sample')

        global path_input
        generate_share_warrants(os.path.join(sample_path, 'sample.xlsx'), path_input, para_list)

        print('股权证生成完毕。')

    B = Button(top, text="生成股权证", command=btn_generator_CallBack)
    B.pack()
    top.mainloop()


from django.db import models

class Visitor(models.Model):
    visitorid = models.IntegerField(db_column='VisitorId', primary_key=True, verbose_name='№')
    login = models.CharField(db_column='Login', max_length=30, blank=True, null=True, verbose_name='Логин')
    passwordhash = models.CharField(db_column='PasswordHash', max_length=64, blank=True, null=True, verbose_name='Хеш пароля')
    session = models.CharField(db_column='Session', max_length=32, blank=True, null=True, verbose_name='Сессионный ключ')
    lastname = models.CharField(db_column='Lastname', max_length=30, blank=True, null=True, verbose_name='Фамилия')
    firstname = models.CharField(db_column='Firstname', max_length=30, blank=True, null=True, verbose_name='Имя')
    patronymic = models.CharField(db_column='Patronymic', max_length=30, blank=True, null=True, verbose_name='Отчество')
    phonenumber = models.CharField(db_column='PhoneNumber', max_length=19, blank=True, null=True, verbose_name='Номер телефона')
    citizenship = models.CharField(db_column='Citizenship', max_length=30, blank=True, null=True, verbose_name='Гражданство')
    passport = models.CharField(db_column='Passport', max_length=15, blank=True, null=True, verbose_name='Серия и номер паспорта')

    class Meta:
        managed = True
        db_table = 'visitor'
        verbose_name = 'Посетитель'
        verbose_name_plural = 'Посетители'

    def __str__(self):
        return self.login

class Roomclass(models.Model):
    roomclassid = models.AutoField(db_column='RoomClassId', primary_key=True, verbose_name='№')
    name = models.CharField(db_column='Name', max_length=30, blank=True, null=True, verbose_name='Название класса')
    cost = models.PositiveIntegerField(db_column='Cost', blank=True, null=True, verbose_name='Стоимость')
    numberofpeople = models.PositiveIntegerField(db_column='NumberOfPeople', blank=True, null=True, verbose_name='Вместимость')
    comment = models.CharField(db_column='Comment', max_length=200, blank=True, null=True, verbose_name='Описание')

    class Meta:
        managed = True
        db_table = 'roomclass'
        verbose_name = 'Класс номера'
        verbose_name_plural = 'Классы номеров'

    def __str__(self):
        return self.name

class Room(models.Model):
    roomid = models.AutoField(db_column='RoomId', primary_key=True, verbose_name='№')
    roomclassid = models.ForeignKey('Roomclass', models.DO_NOTHING, db_column='RoomClassId', blank=True, null=True, verbose_name='Класс номера')
    roomnumber = models.PositiveIntegerField(db_column='RoomNumber', blank=True, null=True, verbose_name='Номер комнаты')
    avaliable = models.BooleanField(db_column='Avaliable', blank=True, null=True, verbose_name='Активный')

    class Meta:
        managed = True
        db_table = 'room'
        verbose_name = 'Номер'
        verbose_name_plural = 'Номера'

    def __str__(self):
        return f"{self.roomnumber} ({self.roomclassid})"

class Orderinfo(models.Model):
    orderid = models.AutoField(db_column='OrderId', primary_key=True, verbose_name='№')
    visitorid = models.ForeignKey('Visitor', models.DO_NOTHING, db_column='VisitorId', blank=True, null=True, verbose_name='Посетитель')
    roomid = models.ForeignKey('Room', models.DO_NOTHING, db_column='RoomId', blank=True, null=True, verbose_name='Комната')
    checkindate = models.DateField(db_column='CheckInDate', blank=True, null=True, verbose_name='Дата заселения')
    checkoutdate = models.DateField(db_column='CheckOutDate', blank=True, null=True, verbose_name='Дата выселения')
    numberofguests = models.PositiveIntegerField(db_column='NumberOfGuests', blank=True, null=True, verbose_name='Количество гостей')
    cost = models.PositiveIntegerField(db_column='Cost', blank=True, null=True, verbose_name='Стоимость')

    class Meta:
        managed = True
        db_table = 'orderinfo'
        verbose_name = 'Заказ'
        verbose_name_plural = 'Заказы'

    def __str__(self):
        return f"Заказ №{self.orderid}"

class Foodtype(models.Model):
    foodtypeid = models.AutoField(db_column='FoodTypeId', primary_key=True, verbose_name='№')
    name = models.CharField(db_column='Name', max_length=30, blank=True, null=True, verbose_name='Название')
    cost = models.PositiveIntegerField(db_column='Cost', blank=True, null=True, verbose_name='Стоимость')
    avaliable = models.BooleanField(db_column='Avaliable', blank=True, null=True, verbose_name='Активный')
    comment = models.CharField(db_column='Comment', max_length=200, blank=True, null=True, verbose_name='Описание')

    class Meta:
        managed = True
        db_table = 'foodtype'
        verbose_name = 'Тип питания'
        verbose_name_plural = 'Типы питания'

    def __str__(self):
        return self.name

class Food(models.Model):
    foodid = models.AutoField(db_column='FoodId', primary_key=True, verbose_name='№')
    orderid = models.ForeignKey('Orderinfo', models.DO_NOTHING, db_column='OrderId', blank=True, null=True, verbose_name='Заказ')
    foodtypeid = models.ForeignKey('Foodtype', models.DO_NOTHING, db_column='FoodTypeId', blank=True, null=True, verbose_name='Тип питания')
    included = models.BooleanField(db_column='Included', blank=True, null=True, verbose_name='Включено в заказ')

    class Meta:
        managed = True
        db_table = 'food'
        verbose_name = 'Питание'
        verbose_name_plural = 'Питание'

class Paymenttype(models.Model):
    paymenttypeid = models.AutoField(db_column='PaymentTypeId', primary_key=True, verbose_name='№')
    name = models.CharField(db_column='Name', max_length=30, blank=True, null=True, verbose_name='Название')
    avaliable = models.BooleanField(db_column='Avaliable', blank=True, null=True, verbose_name='Активный')

    class Meta:
        managed = True
        db_table = 'paymenttype'
        verbose_name = 'Тип оплаты'
        verbose_name_plural = 'Типы оплаты'

    def __str__(self):
        return self.name

class Orderstatus(models.Model):
    orderstatusid = models.AutoField(db_column='OrderStatusId', primary_key=True, verbose_name='№')
    orderid = models.ForeignKey(Orderinfo, models.DO_NOTHING, db_column='OrderId', blank=True, null=True, verbose_name='Заказ')
    paymenttypeid = models.ForeignKey('Paymenttype', models.DO_NOTHING, db_column='PaymentTypeId', blank=True, null=True, verbose_name='Тип оплаты')
    orderactive = models.BooleanField(db_column='OrderActive', blank=True, null=True, verbose_name='Активный')
    orderpayed = models.BooleanField(db_column='OrderPayed', blank=True, null=True, verbose_name='Оплаченный')

    class Meta:
        managed = True
        db_table = 'orderstatus'
        verbose_name = 'Статус заказа'
        verbose_name_plural = 'Статусы заказов'

    def __str__(self):
        return f"Статус заказа №{self.orderid}"

class Addservicetype(models.Model):
    addservicetypeid = models.AutoField(db_column='AddServiceTypeId', primary_key=True, verbose_name='№')
    name = models.CharField(db_column='Name', max_length=30, blank=True, null=True, verbose_name='Название')
    cost = models.PositiveIntegerField(db_column='Cost', blank=True, null=True, verbose_name='Стоимость')
    avaliable = models.BooleanField(db_column='Avaliable', blank=True, null=True, verbose_name='Активный')

    class Meta:
        managed = True
        db_table = 'addservicetype'
        verbose_name = 'Набор доп. услуг'
        verbose_name_plural = 'Наборы доп. услуг'

    def __str__(self):
        return self.name

class Addservices(models.Model):
    addservicesid = models.AutoField(db_column='AddServicesId', primary_key=True, verbose_name='№')
    orderid = models.ForeignKey('Orderinfo', models.DO_NOTHING, db_column='OrderId', blank=True, null=True, verbose_name='Заказ')
    addservicetypeid = models.ForeignKey('Addservicetype', models.DO_NOTHING, db_column='AddServiceTypeId', blank=True, null=True, verbose_name='Тип доп. услуги')

    class Meta:
        managed = True
        db_table = 'addservices'
        verbose_name = 'Доп. услуга'
        verbose_name_plural = 'Доп. услуги'


import torch
import torch.nn as nn
import torchvision
import numpy as np
import torch.utils.model_zoo as model_zoo
import torchvision.models as models
print("PyTorch Version: ", torch.__version__)
print("Torchvision Version: ", torchvision.__version__)

# __all__ = ['ResNet50', 'ResNet101', 'ResNet152','ResNet153']
model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

def Conv1(in_planes, places, stride=2):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_planes, out_channels=places, kernel_size=7, stride=stride, padding=3, bias=False),
        nn.BatchNorm2d(places),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
    )

class Bottleneck(nn.Module):
    def __init__(self, in_places, places, stride=1, downsampling=False, expansion=4):
        super(Bottleneck, self).__init__()
        self.expansion = expansion
        self.downsampling = downsampling

        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_channels=in_places, out_channels=places, kernel_size=1, stride=1, bias=False),
            nn.BatchNorm2d(places),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=places, out_channels=places, kernel_size=3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(places),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=places, out_channels=places * self.expansion, kernel_size=1, stride=1, bias=False),
            nn.BatchNorm2d(places * self.expansion),
        )

        if self.downsampling:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels=in_places, out_channels=places * self.expansion, kernel_size=1, stride=stride,
                          bias=False),
                nn.BatchNorm2d(places * self.expansion)
            )
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        residual = x
        out = self.bottleneck(x)

        if self.downsampling:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)
        return out

class _PositionAttentionModule(nn.Module):
    """ Position attention module"""

    def __init__(self, in_channels, **kwargs):
        super(_PositionAttentionModule, self).__init__()
        self.conv_b = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.conv_c = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.conv_d = nn.Conv2d(in_channels, in_channels, 1)
        self.alpha = nn.Parameter(torch.zeros(1))
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        batch_size, _, height, width = x.size()  # [b,512,14,14]
        feat_b = self.conv_b(x).view(batch_size, -1, height * width).permute(0, 2, 1)
        feat_c = self.conv_c(x).view(batch_size, -1, height * width)
        attention_s = self.softmax(torch.bmm(feat_b, feat_c))
        feat_d = self.conv_d(x).view(batch_size, -1, height * width)
        feat_e = torch.bmm(feat_d, attention_s.permute(0, 2, 1)).view(batch_size, -1, height, width)
        out = self.alpha * feat_e + x  # [b,512,14,14]

        return out

class _ChannelAttentionModule(nn.Module):
    """Channel attention module"""

    def __init__(self, **kwargs):
        super(_ChannelAttentionModule, self).__init__()
        self.beta = nn.Parameter(torch.zeros(1))
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        batch_size, _, height, width = x.size()
        feat_a = x.view(batch_size, -1, height * width)
        feat_a_transpose = x.view(batch_size, -1, height * width).permute(0, 2, 1)
        attention = torch.bmm(feat_a, feat_a_transpose)
        attention_new = torch.max(attention, dim=-1, keepdim=True)[0].expand_as(attention) - attention
        attention = self.softmax(attention_new)

        feat_e = torch.bmm(attention, feat_a).view(batch_size, -1, height, width)
        out = self.beta * feat_e + x  #

        return out

class ResNet(nn.Module):
    def __init__(self, blocks, norm_layer=nn.BatchNorm2d, norm_kwargs=None, num_classes=1000, expansion=4, **kwargs):
        super(ResNet, self).__init__()
        self.expansion = expansion

        self.conv1 = Conv1(in_planes=3, places=64)

        self.layer1 = self.make_layer(in_places=64, places=64, block=blocks[0], stride=1)

        self.layer1_down = nn.Sequential(nn.Conv2d(256, 1024, 1, stride=2, bias=True),
                                         norm_layer(1024, **({} if norm_kwargs is None else norm_kwargs)),
                                         nn.Conv2d(1024, 2048, 1, stride=2, bias=True),
                                         norm_layer(2048, **({} if norm_kwargs is None else norm_kwargs)))

        self.layer2 = self.make_layer(in_places=256, places=128, block=blocks[1], stride=2)

        self.conv_p1 = nn.Sequential(
            nn.Conv2d(512, 512 // 4, 3, padding=1, bias=False),
            norm_layer(512 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c1 = nn.Sequential(
            nn.Conv2d(512, 512 // 4, 3, padding=1, bias=False),
            norm_layer(512 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.pam = _PositionAttentionModule(512 // 4, **kwargs)
        self.cam = _ChannelAttentionModule(**kwargs)
        self.conv_p2 = nn.Sequential(
            nn.Conv2d(512 // 4, 512 // 4, 3, padding=1, bias=True),
            norm_layer(512 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c2 = nn.Sequential(
            nn.Conv2d(512 // 4, 512 // 4, 3, padding=1, bias=True),
            norm_layer(512 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True))

        self.fusion_cp1 = nn.Sequential(nn.Conv2d(512 // 4, 512, 1, bias=True),
                                        nn.BatchNorm2d(512))

        self.down = nn.Sequential(nn.Conv2d(512, 1024, 1, stride=2, bias=True), nn.BatchNorm2d(1024))

        self.layer3 = self.make_layer(in_places=512, places=256, block=blocks[2], stride=2)

        self.conv_p1_layer3 = nn.Sequential(
            nn.Conv2d(1024, 1024 // 4, 3, padding=1, bias=False),
            norm_layer(1024 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c1_layer3 = nn.Sequential(
            nn.Conv2d(1024, 1024 // 4, 3, padding=1, bias=False),
            norm_layer(1024 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.pam_layer3 = _PositionAttentionModule(1024 // 4, **kwargs)
        self.cam_layer3 = _ChannelAttentionModule(**kwargs)
        self.conv_p2_layer3 = nn.Sequential(
            nn.Conv2d(1024 // 4, 1024 // 4, 3, padding=1, bias=True),
            norm_layer(1024 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c2_layer3 = nn.Sequential(
            nn.Conv2d(1024 // 4, 1024 // 4, 3, padding=1, bias=True),
            norm_layer(1024 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True))

        self.fusion_cp1_layer3 = nn.Sequential(nn.Conv2d(1024 // 4, 1024, 1, bias=True),
                                               nn.BatchNorm2d(1024))

        self.down_layer3 = nn.Sequential(nn.Conv2d(1024, 2048, 1, stride=1, bias=True), nn.BatchNorm2d(2048))

        self.layer4 = self.make_layer(in_places=1024, places=512, block=blocks[3], stride=1)

        self.conv_p1_layer4 = nn.Sequential(
            nn.Conv2d(2048, 2048 // 4, 3, padding=1, bias=False),
            norm_layer(2048 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c1_layer4 = nn.Sequential(
            nn.Conv2d(2048, 2048 // 4, 3, padding=1, bias=False),
            norm_layer(2048 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.pam_layer4 = _PositionAttentionModule(2048 // 4, **kwargs)
        self.cam_layer4 = _ChannelAttentionModule(**kwargs)
        self.conv_p2_layer4 = nn.Sequential(
            nn.Conv2d(2048 // 4, 2048 // 4, 3, padding=1, bias=True),
            norm_layer(2048 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True)
        )
        self.conv_c2_layer4 = nn.Sequential(
            nn.Conv2d(2048 // 4, 2048 // 4, 3, padding=1, bias=True),
            norm_layer(2048 // 4, **({} if norm_kwargs is None else norm_kwargs)),
            nn.ReLU(True))

        self.fusion_cp1_layer4 = nn.Sequential(nn.Conv2d(2048 // 4, 2048, 1, bias=True),
                                               nn.BatchNorm2d(2048))
        self.avgpool = nn.AvgPool2d(14, stride=1)

        self.fc = nn.Linear(2048,num_classes)

        # for m in self.modules():
        #     if isinstance(m, nn.Conv2d):
        #         nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        #     elif isinstance(m, nn.BatchNorm2d):
        #         nn.init.constant_(m.weight, 1)
        #         nn.init.constant_(m.bias, 0)

    def make_layer(self, in_places, places, block, stride):
        layers = []
        layers.append(Bottleneck(in_places, places, stride, downsampling=True))
        for i in range(1, block):
            layers.append(Bottleneck(places * self.expansion, places))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)

        x = self.layer1(x)  # [b,256,56,56]
        x_layer1 = self.layer1_down(x)  # torch.Size([1, 2048, 14, 14])
        layer2 = self.layer2(x)  # [b,512,28,28]

        feat_p = self.conv_p1(layer2)  # torch.Size([1, 128, 28, 28])
        feat_p = self.pam(feat_p)  # torch.Size([1, 128, 28, 28])
        feat_p = self.conv_p2(feat_p)  # torch.Size([1, 128, 28, 28])

        feat_c = self.conv_c1(layer2)  # torch.Size([1, 128, 28, 28])

        feat_c = self.cam(feat_c)  # torch.Size([1, 128, 28, 28])

        feat_c = self.conv_c2(feat_c)  # torch.Size([1, 128, 28, 28])

        feat_fusion_layer1 = feat_p + feat_c  # torch.Size([1, 128, 28, 28])
        x = self.fusion_cp1(feat_fusion_layer1)  # torch.Size([1, 512, 28, 28])

        # print("new_F2生成完毕")

        x_downsample = self.down(x)  # torch.Size([1, 1024, 14, 14])

        layer3 = self.layer3(x)  # torch.Size([1, 1024, 14, 14])
        layer3 = x_downsample + layer3
        feat_p_layer3 = self.conv_p1_layer3(layer3)
        feat_p_layer3 = self.pam_layer3(feat_p_layer3)
        feat_p_layer3 = self.conv_p2_layer3(feat_p_layer3)

        feat_c_layer3 = self.conv_c1_layer3(layer3)

        feat_c_layer3 = self.cam_layer3(feat_c_layer3)

        feat_c_layer3 = self.conv_c2_layer3(feat_c_layer3)
        feat_fusion_layer3 = feat_p_layer3 + feat_c_layer3
        x = self.fusion_cp1_layer3(feat_fusion_layer3)  # torch.Size([1, 1024, 14, 14])

        # print("new_F3生成完毕")

        layer4 = self.layer4(x) + self.down_layer3(x)

        feat_p_layer4 = self.conv_p1_layer4(layer4)
        feat_p_layer4 = self.pam_layer4(feat_p_layer4)
        feat_p_layer4 = self.conv_p2_layer4(feat_p_layer4)

        feat_c_layer4 = self.conv_c1_layer4(layer4)

        feat_c_layer4 = self.cam_layer4(feat_c_layer4)

        feat_c_layer4 = self.conv_c2_layer4(feat_c_layer4)
        feat_fusion_layer4 = feat_p_layer4 + feat_c_layer4
        x = self.fusion_cp1_layer4(feat_fusion_layer4)  # torch.Size([1, 2048, 14, 14])
        x = x_layer1 + x
        # print("new_F4生成完毕")

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)

        x = self.fc(x)
        return x

def ResNet50():
    return ResNet([3, 4, 6, 3])

def ResNet101():
    return ResNet([3, 4, 23, 3])

def ResNet153():
    ResNet153 = models.resnet152(pretrained=True)
    return ResNet153

def ResNet152():

    resNet152 =models.resnet152(pretrained=True)
    pretrained_dict = resNet152.state_dict()

    model = ResNet([3, 8, 36, 3])
    model_dict = model.state_dict()

    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    # 更新现有的model_dict
    model_dict.update(pretrained_dict)

    # 加载真正需要的state_dict
    model.load_state_dict(model_dict)
    return model

if __name__ == '__main__':
    # model = torchvision.models.resnet50()
    model = ResNet152()
    #  print(model)

    input = torch.randn(1, 3, 224, 224)
    out = model(input)
    print(out.shape)


from PySide2.QtCore import *
from PySide2.QtGui import *
from PySide2.QtWidgets import *

from swatchM8.core.util.Helpers import convertToQColor

from swatchM8.ui.qt.widgets.ColorPickerWidget import ColorPickerWidget

from PySide2 import QtCore, QtGui, QtWidgets

class AddPartWidget(QWidget):

    def __init__(self, MainUI, Main):
        # type: (object, object) -> None

        super(AddPartWidget, self).__init__(None)  # Calls the init of the super class -> QWidget

        self.setupUi()
        self.connectElements()
        self.connectInputChecks()
        self.MainUI = MainUI
        self.Main = Main

        layout = QHBoxLayout()
        layout.setMargin(0)
        layout.addWidget(self.AddPartFrame)
        self.setLayout(layout)

        self.ColorPicker = ColorPickerWidget()
        self.ColorPickerPlaceHolder.setLayout(self.ColorPicker)

        self.ModelComboBox.setCurrentIndex(0)

        self.MainUI.setInputInvalid(self.ModelNameLineEdit)
        self.MainUI.setInputInvalid(self.PartNameLineEdit)

        self.UniqueSwatchRadioButton.setChecked(True)
        self.setUniqueSwatchSetup()

    def setupUi(self):
        self.AddPartFrame = QtWidgets.QFrame()
        self.AddPartFrame.setGeometry(QtCore.QRect(0, 0, 311, 261))
        self.AddPartFrame.setFrameShape(QtWidgets.QFrame.StyledPanel)
        self.AddPartFrame.setFrameShadow(QtWidgets.QFrame.Raised)
        self.AddPartFrame.setObjectName("AddPartFrame")
        self.AddPartLabel = QtWidgets.QLabel(self.AddPartFrame)
        self.AddPartLabel.setGeometry(QtCore.QRect(10, 10, 141, 21))
        font = QtGui.QFont()
        font.setPointSize(16)
        self.AddPartLabel.setFont(font)
        self.AddPartLabel.setObjectName("AddPartLabel")
        self.CreateButton = QtWidgets.QPushButton(self.AddPartFrame)
        self.CreateButton.setGeometry(QtCore.QRect(200, 180, 81, 71))
        self.CreateButton.setObjectName("CreateButton")
        self.ColorPickerPlaceHolder = QtWidgets.QFrame(self.AddPartFrame)
        self.ColorPickerPlaceHolder.setGeometry(QtCore.QRect(10, 170, 181, 81))
        self.ColorPickerPlaceHolder.setFrameShape(QtWidgets.QFrame.StyledPanel)
        self.ColorPickerPlaceHolder.setFrameShadow(QtWidgets.QFrame.Raised)
        self.ColorPickerPlaceHolder.setObjectName("ColorPickerPlaceHolder")
        self.horizontalLayoutWidget_2 = QtWidgets.QWidget(self.AddPartFrame)
        self.horizontalLayoutWidget_2.setGeometry(QtCore.QRect(10, 40, 271, 31))
        self.horizontalLayoutWidget_2.setObjectName("horizontalLayoutWidget_2")
        self.ModelLayout = QtWidgets.QHBoxLayout(self.horizontalLayoutWidget_2)
        self.ModelLayout.setContentsMargins(-1, 2, -1, 2)
        self.ModelLayout.setObjectName("ModelLayout")
        self.ModelLabel = QtWidgets.QLabel(self.horizontalLayoutWidget_2)
        self.ModelLabel.setObjectName("ModelLabel")
        self.ModelLayout.addWidget(self.ModelLabel)
        self.ModelStackedWidget = QtWidgets.QStackedWidget(self.horizontalLayoutWidget_2)
        self.ModelStackedWidget.setToolTipDuration(0)
        self.ModelStackedWidget.setObjectName("ModelStackedWidget")
        self.page_5 = QtWidgets.QWidget()
        self.page_5.setObjectName("page_5")
        self.ModelComboBox = QtWidgets.QComboBox(self.page_5)
        self.ModelComboBox.setGeometry(QtCore.QRect(18, 3, 141, 21))
        self.ModelComboBox.setObjectName("ModelComboBox")
        self.NewModelButton = QtWidgets.QPushButton(self.page_5)
        self.NewModelButton.setGeometry(QtCore.QRect(160, 2, 51, 23))
        self.NewModelButton.setObjectName("NewModelButton")
        self.ModelStackedWidget.addWidget(self.page_5)
        self.page_6 = QtWidgets.QWidget()
        self.page_6.setObjectName("page_6")
        self.ModelNameLineEdit = QtWidgets.QLineEdit(self.page_6)
        self.ModelNameLineEdit.setGeometry(QtCore.QRect(20, 3, 141, 21))
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(self.ModelNameLineEdit.sizePolicy().hasHeightForWidth())
        self.ModelNameLineEdit.setSizePolicy(sizePolicy)
        self.ModelNameLineEdit.setMaximumSize(QtCore.QSize(240, 16777215))
        self.ModelNameLineEdit.setReadOnly(False)
        self.ModelNameLineEdit.setObjectName("ModelNameLineEdit")
        self.ExistingModelButton = QtWidgets.QPushButton(self.page_6)
        self.ExistingModelButton.setGeometry(QtCore.QRect(160, 2, 61, 23))
        self.ExistingModelButton.setObjectName("ExistingModelButton")
        self.ModelStackedWidget.addWidget(self.page_6)
        self.ModelLayout.addWidget(self.ModelStackedWidget)
        self.horizontalLayoutWidget_3 = QtWidgets.QWidget(self.AddPartFrame)
        self.horizontalLayoutWidget_3.setGeometry(QtCore.QRect(10, 80, 271, 23))
        self.horizontalLayoutWidget_3.setObjectName("horizontalLayoutWidget_3")
        self.PartLayout = QtWidgets.QHBoxLayout(self.horizontalLayoutWidget_3)
        self.PartLayout.setContentsMargins(-1, 2, 2, 2)
        self.PartLayout.setObjectName("PartLayout")
        self.PartNameLabel = QtWidgets.QLabel(self.horizontalLayoutWidget_3)
        self.PartNameLabel.setObjectName("PartNameLabel")
        self.PartLayout.addWidget(self.PartNameLabel)
        self.PartNameLineEdit = QtWidgets.QLineEdit(self.horizontalLayoutWidget_3)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(self.PartNameLineEdit.sizePolicy().hasHeightForWidth())
        self.PartNameLineEdit.setSizePolicy(sizePolicy)
        self.PartNameLineEdit.setMaximumSize(QtCore.QSize(240, 16777215))
        self.PartNameLineEdit.setReadOnly(False)
        self.PartNameLineEdit.setObjectName("PartNameLineEdit")
        self.PartLayout.addWidget(self.PartNameLineEdit)
        self.horizontalLayoutWidget = QtWidgets.QWidget(self.AddPartFrame)
        self.horizontalLayoutWidget.setGeometry(QtCore.QRect(10, 140, 271, 31))
        self.horizontalLayoutWidget.setObjectName("horizontalLayoutWidget")
        self.NameLayout = QtWidgets.QHBoxLayout(self.horizontalLayoutWidget)
        self.NameLayout.setContentsMargins(-1, 2, -1, 2)
        self.NameLayout.setObjectName("NameLayout")
        self.SwatchLabel = QtWidgets.QLabel(self.horizontalLayoutWidget)
        self.SwatchLabel.setObjectName("SwatchLabel")
        self.NameLayout.addWidget(self.SwatchLabel)
        self.SwatchStackedWidget = QtWidgets.QStackedWidget(self.horizontalLayoutWidget)
        self.SwatchStackedWidget.setObjectName("SwatchStackedWidget")
        self.page_3 = QtWidgets.QWidget()
        self.page_3.setObjectName("page_3")
        self.SwatchComboBox = QtWidgets.QComboBox(self.page_3)
        self.SwatchComboBox.setGeometry(QtCore.QRect(0, 3, 171, 21))
        self.SwatchComboBox.setObjectName("SwatchComboBox")
        self.NewSwatchButton = QtWidgets.QPushButton(self.page_3)
        self.NewSwatchButton.setGeometry(QtCore.QRect(171, 2, 51, 23))
        self.NewSwatchButton.setObjectName("NewSwatchButton")
        self.SwatchStackedWidget.addWidget(self.page_3)
        self.page_4 = QtWidgets.QWidget()
        self.page_4.setObjectName("page_4")
        self.SwatchNameLineEdit = QtWidgets.QLineEdit(self.page_4)
        self.SwatchNameLineEdit.setGeometry(QtCore.QRect(0, 3, 161, 21))
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(self.SwatchNameLineEdit.sizePolicy().hasHeightForWidth())
        self.SwatchNameLineEdit.setSizePolicy(sizePolicy)
        self.SwatchNameLineEdit.setMaximumSize(QtCore.QSize(240, 16777215))
        self.SwatchNameLineEdit.setReadOnly(False)
        self.SwatchNameLineEdit.setObjectName("SwatchNameLineEdit")
        self.ExistingSwatchButton = QtWidgets.QPushButton(self.page_4)
        self.ExistingSwatchButton.setGeometry(QtCore.QRect(162, 2, 51, 23))
        self.ExistingSwatchButton.setObjectName("ExistingSwatchButton")
        self.SwatchStackedWidget.addWidget(self.page_4)
        self.NameLayout.addWidget(self.SwatchStackedWidget)
        self.horizontalLayoutWidget_4 = QtWidgets.QWidget(self.AddPartFrame)
        self.horizontalLayoutWidget_4.setGeometry(QtCore.QRect(10, 120, 211, 21))
        self.horizontalLayoutWidget_4.setObjectName("horizontalLayoutWidget_4")
        self.RadioButtonLayout = QtWidgets.QHBoxLayout(self.horizontalLayoutWidget_4)
        self.RadioButtonLayout.setContentsMargins(0, 0, 0, 0)
        self.RadioButtonLayout.setObjectName("RadioButtonLayout")
        self.UniqueSwatchRadioButton = QtWidgets.QRadioButton(self.horizontalLayoutWidget_4)
        self.UniqueSwatchRadioButton.setObjectName("UniqueSwatchRadioButton")
        self.RadioButtonLayout.addWidget(self.UniqueSwatchRadioButton)
        self.SharedSwatchRadioButton = QtWidgets.QRadioButton(self.horizontalLayoutWidget_4)
        self.SharedSwatchRadioButton.setObjectName("SharedSwatchRadioButton")
        self.RadioButtonLayout.addWidget(self.SharedSwatchRadioButton)
        self.line = QtWidgets.QFrame(self.AddPartFrame)
        self.line.setGeometry(QtCore.QRect(10, 108, 271, 8))
        self.line.setFrameShape(QtWidgets.QFrame.HLine)
        self.line.setFrameShadow(QtWidgets.QFrame.Sunken)
        self.line.setObjectName("line")

        self.retranslateUi()
        self.ModelStackedWidget.setCurrentIndex(0)
        self.SwatchStackedWidget.setCurrentIndex(0)

    def retranslateUi(self):
        self.AddPartLabel.setText(QtWidgets.QApplication.translate("FormName", "Add Part", None, -1))
        self.CreateButton.setText(QtWidgets.QApplication.translate("FormName", "Create", None, -1))
        self.ModelLabel.setText(QtWidgets.QApplication.translate("FormName", "Model:", None, -1))
        self.NewModelButton.setText(QtWidgets.QApplication.translate("FormName", "New", None, -1))
        self.ExistingModelButton.setText(QtWidgets.QApplication.translate("FormName", "Pick", None, -1))
        self.PartNameLabel.setText(QtWidgets.QApplication.translate("FormName", "PartName:", None, -1))
        self.SwatchLabel.setText(QtWidgets.QApplication.translate("FormName", "Swatch:", None, -1))
        self.NewSwatchButton.setText(QtWidgets.QApplication.translate("FormName", "New", None, -1))
        self.ExistingSwatchButton.setText(QtWidgets.QApplication.translate("FormName", "Pick", None, -1))
        self.UniqueSwatchRadioButton.setText(QtWidgets.QApplication.translate("FormName", "Unique Swatch", None, -1))
        self.SharedSwatchRadioButton.setText(QtWidgets.QApplication.translate("FormName", "Shared Swatch", None, -1))

    def connectElements(self):
        # type: () -> None

        self.NewSwatchButton.clicked.connect(self.newSwatchButton)
        self.NewModelButton.clicked.connect(self.newModelButton)
        self.ExistingSwatchButton.clicked.connect(self.existingSwatchButton)
        self.ExistingModelButton.clicked.connect(self.existingModelButton)

        self.SwatchComboBox.currentIndexChanged.connect(self.updateSwatch)

        self.CreateButton.clicked.connect(self.addButton)
        self.UniqueSwatchRadioButton.toggled.connect(self.setUniqueSwatchSetup)
        self.SharedSwatchRadioButton.toggled.connect(self.setSharedSwatchSetup)

        self.ModelComboBox.currentIndexChanged.connect(self.generateUniqueSwatchName)
        self.ModelNameLineEdit.textEdited.connect(self.generateUniqueSwatchName)
        self.PartNameLineEdit.textEdited.connect(self.generateUniqueSwatchName)

    def connectInputChecks(self):

        self.ModelComboBox.currentIndexChanged.connect(lambda: self.MainUI.checkIfPartNameEntryIsValid(self.PartNameLineEdit, self.ModelComboBox.currentData()))
        self.ModelNameLineEdit.textChanged.connect(lambda: self.MainUI.checkIfModelNameEntryIsValid(self.ModelNameLineEdit))
        self.PartNameLineEdit.textChanged.connect(self.inputCheckPartName)
        self.SwatchNameLineEdit.textChanged.connect(self.inputCheckSwatchName)

    def inputCheckPartName(self):

        if self.PartNameLineEdit.text() == "":
            self.MainUI.setInputInvalid(self.PartNameLineEdit)
            return

        if self.ModelStackedWidget.currentIndex() == 1:  # If new Model aka no parts yet
            self.MainUI.setInputValid(self.PartNameLineEdit)
            return

        self.MainUI.checkIfPartNameEntryIsValid(self.PartNameLineEdit, self.ModelComboBox.currentData())

    def inputCheckSwatchName(self):

        if self.SwatchStackedWidget.currentIndex() == 0:  # als er een gepicked word, geen check
            return

        if self.SwatchNameLineEdit.text() == "" and self.SharedSwatchRadioButton.isChecked():  # Leeg en customName, sws niet valid
            self.MainUI.setInputInvalid(self.SwatchNameLineEdit)
            return

        if self.Main.texture == None:  # Bij init code geeft hij anders een Attribute error
            return
        self.MainUI.checkifSwatchNameIsValid(self.SwatchNameLineEdit)

    def setNewData(self):
        # type: () -> None

        self.updateModelDropdown()

        if self.ModelComboBox.count() == 0:
            self.newModelButton()

        self.updateSwatchDropdown()

        self.ColorPicker.setNewData(QColor('white'))

        # self.updateSwatch(self.SwatchComboBox.currentIndex())

    def updateSwatchDropdown(self):
        # type: () -> None

        self.SwatchComboBox.clear()
        for index, item in enumerate(self.Main.texture.swatches):
            self.SwatchComboBox.addItem(item.name)
            self.SwatchComboBox.setItemData(index, item)

    def updateModelDropdown(self):
        # type: () -> None

        self.ModelComboBox.clear()
        for index, item in enumerate(self.Main.models):
            self.ModelComboBox.addItem(item.name)
            self.ModelComboBox.setItemData(index, item)

    def updateSwatch(self, index):
        # type: (int) -> None

        selectedSwatch = self.SwatchComboBox.itemData(index)
        if selectedSwatch != None:
            self.ColorPicker.setNewData(convertToQColor(selectedSwatch.color))
            # self.ColorPicker.setNewData(QColor('white'))

    def newSwatchButton(self):
        # type: () -> None

        self.SwatchStackedWidget.setCurrentIndex(1)
        self.ColorPicker.setNewData(QColor('white'))
        self.ColorPickerPlaceHolder.setEnabled(True)

    def newModelButton(self):
        # type: () -> None

        self.ModelStackedWidget.setCurrentIndex(1)

    def setUniqueSwatchSetup(self):
        # type: () -> None

        self.SwatchStackedWidget.setCurrentIndex(1)
        self.ExistingSwatchButton.setEnabled(0)
        self.SwatchNameLineEdit.setEnabled(0)
        self.ColorPickerPlaceHolder.setEnabled(True)
        self.ColorPicker.setNewData(QColor('white'))
        self.generateUniqueSwatchName()

    def setSharedSwatchSetup(self):
        # type: () -> None

        self.ExistingSwatchButton.setEnabled(1)
        self.SwatchNameLineEdit.setEnabled(1)
        self.SwatchNameLineEdit.setText('')

    def generateUniqueSwatchName(self):
        # type: () -> None

        if self.UniqueSwatchRadioButton.isChecked():

            if self.ModelStackedWidget.currentIndex() == 0:

                modelName = self.ModelComboBox.currentText()

            else:
                modelName = self.ModelNameLineEdit.text()

            partName = self.PartNameLineEdit.text()

            if partName == '':
                partName = 'PartName'

            if modelName == '':
                modelName = 'ModelName'

            self.SwatchNameLineEdit.setText(modelName + "-" + partName)

    def existingSwatchButton(self):
        # type: () -> None

        self.updateSwatch(self.SwatchComboBox.currentIndex())
        self.SwatchStackedWidget.setCurrentIndex(0)
        self.ColorPickerPlaceHolder.setEnabled(False)

    def existingModelButton(self):
        # type: () -> None

        self.ModelStackedWidget.setCurrentIndex(0)

    def addButton(self):
        # type: () -> None

        color = self.ColorPicker.getPickedColor()

        if self.ModelStackedWidget.currentIndex() == 0:
            model = self.ModelComboBox.currentData()
        else:
            model = self.Main.createModel(self.ModelNameLineEdit.text())

        if self.SwatchStackedWidget.currentIndex() == 0:
            swatch = self.SwatchComboBox.itemData(self.SwatchComboBox.currentIndex())
        else:
            swatch = self.Main.texture.createSwatch(self.SwatchNameLineEdit.text(), color)

        # TODO: model verwijderen als het fout loopt bij de swatch

        self.Main.addPartToSwatch(swatch, model, self.PartNameLineEdit.text())
        self.Main.app.unwrapSelection(swatch.coordinate, self.Main.texture.amountOfSwatchesOnRow)

        self.MainUI.applyChanges()
        self.reset()

    def reset(self):
        # type: () -> None

        self.SwatchNameLineEdit.setText('')
        self.ModelNameLineEdit.setText('')
        self.PartNameLineEdit.setText('')
        self.ModelStackedWidget.setCurrentIndex(0)

        self.setNewData()

    def setComboBoxModel(self, model):
        # if the add part button is pressed trought the edit model frame, the correct model will be set

        index = self.ModelComboBox.findData(model)
        self.ModelComboBox.setCurrentIndex(index)


# @copy by sobhan siamak

import numpy as np
import pandas as pd
import  matplotlib.pyplot as plt
import time
from datetime import datetime
import networkx as nx
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import classification_report
# from sklearn import cross_validation
from sklearn.model_selection import cross_val_score,KFold, cross_validate
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_curve
# import scikitplot as skplt

start_time = datetime.now()

db = pd.read_csv("ppi_dmelanogaster.csv")
dbmain = pd.read_csv("comppi--compartments--tax_dmelanogaster_loc_all.txt", sep='\t')
labels = pd.read_csv("Labels.csv")
labels = labels[['locus','essentiality status']]
print(labels.head())

# la1 = pd.DataFrame(la)
# lb1 = pd.DataFrame(lb)
# la1.to_csv("goa.csv" , index=False)
# lb1.to_csv("gob.csv", index=False)
goa = dbmain.iloc[0:,3]
gob = dbmain.iloc[0:,11]

la = [x.split("|") for x in goa]
lb = [x.split("|") for x in gob]

alpha = 1 #coif of gene ontology
beta = 1 #coif of nieghborhood
gama = 0.001 #coif of score

print("la is:", la[0])
# print(goa)
print("lb is :", lb[0])
intersect = np.intersect1d(la[0],lb[0])
print(intersect)

go = []
for i in range(len(la)):
  go.append(len(np.intersect1d(la[i], lb[i]))/ len(la[i]+lb[i]))

######Change weigth #1

print(db.head())
db['go'] = go
db['Interaction Score'] = gama*db['Interaction Score'] + alpha*db['go']
db = db.drop(['go'], axis=1)
print(db.head())

db2 = np.array(db)
Pra = db.iloc[1:,0]
Prb = db.iloc[1:,1]
Sc = db.iloc[1:,2]
prota = db2[0:,0]
protb = db2[0:,1]
score = db2[0:,2]
Pa = pd.read_csv("ProtA.tab", sep='\t')
Pb = pd.read_csv("ProtB.tab", sep='\t')
frames =[Pa,Pb]
proteins = (pd.concat(frames))
proteins = proteins.drop_duplicates()
print(proteins.head())

#Labeling Dataset
prlabels = []
prot = proteins.iloc[:,1]

for i in prot:
    if(sum(labels['locus'].isin([i])) == 0):
        prlabels.append(['unknown'])
        continue
    temp = labels[labels['locus'].isin([i])]
    lbl = list(temp['essentiality status'])
    prlabels.append(lbl)

for i in range(len(prlabels)):
    if prlabels[i] == ['NE,NE']:
        prlabels[i] = ['NE']
    if prlabels[i] == ['E,NE']:
        prlabels[i] = ['E']
    if prlabels[i] == ['NE,E']:
        prlabels[i] = ['E']
    if prlabels[i] == ['E,E']:
        prlabels[i] = ['E']
    if prlabels[i] == []:
        prlabels[i] = ['unknown']

print(prlabels)
print(len(prlabels))
proteins['Labels'] = np.squeeze(prlabels)
print(proteins.head())
FProteins = proteins.drop('To', axis=1)
FProteins = FProteins.rename(columns={"From": "Proteins"})
FProteins = FProteins.drop_duplicates()
print(FProteins.head())
print(FProteins)
fpcsv = pd.DataFrame(FProteins)
fpcsv.to_csv("FPR.csv", index=False)

# unknown = sum(proteins['Labels'].isin(['unknown']))
# essential = sum(proteins['Labels'].isin(['E']))
# nessential = sum(proteins['Labels'].isin(['NE']))

unknown = sum(FProteins['Labels'].isin(['unknown']))
essential = sum(FProteins['Labels'].isin(['E']))
nessential = sum(FProteins['Labels'].isin(['NE']))

print("the number of unknown of proteins are:", unknown)
print("the number of essential of proteins are:", essential)
print("the number of not essential of proteins are:", nessential)

#Creation Graph
G2 = nx.from_pandas_edgelist(db, "Interactor A", "Interactor B", ['Interaction Score'], create_using= nx.Graph())
print(G2.nodes())
print(len(G2.nodes()))

nodes = list(G2.nodes())
nodelabels = []
for i in nodes:
    tmp = FProteins[FProteins['Proteins'].isin([i])]
    lbls = list(tmp['Labels'])
    nodelabels.append(lbls)

print("length nodelabels is:", len(nodelabels))
print(G2.nodes())
print(len(G2.nodes()))
Fnodelabels = []
for i in range(len(nodelabels)):
    if (nodelabels[i] == ['unknown']):
        G2.remove_node(nodes[i])
    else:
        Fnodelabels.append(nodelabels[i])
print(G2.nodes())

FNlabels =[]
for i in Fnodelabels:
    if i == ['NE']:
        FNlabels.append(0)
    else:FNlabels.append(1)
print(FNlabels)
print(len(FNlabels))

print("length Fnodelabels is:", len(Fnodelabels))
print("length G2.nodes is:", len(G2.nodes()))
Fnodes = G2.nodes()

######Centrality
central1 = nx.eigenvector_centrality(G2,weight = 'Interaction Score', max_iter=1000)
central1 = list(central1.values())
# central11 = pd.DataFrame(central1)
# central11.to_csv("central1.csv", index=False)
print(central1)
central2 = nx.degree_centrality(G2)
central2 = list(central2.values())
# central22 = pd.DataFrame(central2)
# central22.to_csv("central2.csv", index=False)
print(central2)
# central2 = nx.degree_pearson_correlation_coefficient(G2,x='P26802', y='Q9W5A9')
# print("pearson corralation is:", central2)
# print(type(central1))
#####High Time Complexity
# central3 = nx.closeness_centrality(G2)
# central3 = list(central3.values())
# central33 = pd.DataFrame(central3)
# central33.to_csv("central3.csv", index=False)
# print(central3)

#####High Time Complexity
# central4 = nx.betweenness_centrality(G2)
# central4 = list(central4.values())
# central44 = pd.DataFrame(central4)
# central44.to_csv("central4.csv", index=False)
# print(central4)

#####High Time Complexity
# central5 = nx.subgraph_centrality(G2)
# central5 = list(central5.values())
# central55 = pd.DataFrame(central5)
# central55.to_csv("central5.csv", index=False)
# print(central5)

###########Classifiers
# DecisionTreeClassifier(max_depth=5),\
# RandomForestClassifier(max_depth=5, n_estimators=10, max_features=1),

# central7 = nx.clustering(G2)
# print("clustering is :", central7)

print(len(Fnodes))
print(len(central1))
print(len(Fnodelabels))
FinalP = pd.DataFrame()
FinalP['Proteins'] = Fnodes
FinalP['C1'] = central1
FinalP['C2'] = central2
c3 = pd.read_csv("central3.csv")
c4 = pd.read_csv("central4.csv")
c5 = pd.read_csv("central5.csv")
print(c3.head())
print(c4.head())
print(c5.head())
FinalP['C3'] = c3
FinalP['C4'] = c4
FinalP['C5'] = c5
#labeling

FinalP['Labels'] = FNlabels
print(FinalP.head())
FinalP.to_csv("Finalp.csv", index=False)

#########################################################################
#Classification

FPro = pd.read_csv("Finalp.csv", header=None)
print(FPro.head())
x = FPro.iloc[1:,1:-1]
y = FPro.iloc[1:,-1]
target_names = ['0', '1']
# print(x.iloc[:,0])
# print(y)
# print(len(y))
X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.20)
# clf = RandomForestClassifier(n_estimators=1000, max_depth=4,random_state=5,class_weight ='imbalanced')
# clf.fit(X_train, y_train)
# result = clf.predict(X_test)
# print(classification_report(y_test, result,target_names=target_names))

clf = KNeighborsClassifier(n_neighbors=5)
clf.fit(X_train, y_train)
result2 = clf.predict(X_test)
print(classification_report(y_test, result2, target_names=target_names))

clf1 = SVC(gamma=2, C=1)
# clf1 = MLPClassifier(alpha=1, max_iter=1000)
# clf1 = GaussianNB()
clf1.fit(X_train, y_train)
result3 = clf1.predict(X_test)
print(classification_report(y_test, result3, target_names=target_names))

clf2 = LogisticRegression()
clf2.fit(X_train,y_train)
result4 = clf2.predict(X_test)
print(classification_report(y_test, result4, target_names=target_names))

clf4 = MLPClassifier(alpha=0.1, max_iter=1000)
clf4.fit(X_train,y_train)
result5 = clf4.predict(X_test)
print(classification_report(y_test, result5, target_names=target_names))

clf5 = GaussianNB()
clf5.fit(X_train,y_train)
result6 = clf5.predict(X_test)
print(classification_report(y_test, result6, target_names=target_names))

# score = cross_val_score(clf, x, y, cv=5)
# print(score)
######KNN
scoring = ['precision_macro', 'recall_macro', 'accuracy']
scores2 = cross_validate(clf, x, y, cv=5, scoring=scoring)
print("recall for 5-fold cross validation is:")
print(scores2['test_recall_macro'])
yknn1 = scores2['test_recall_macro']
print("precision for 5-fold cross validation is:")
print(scores2['test_precision_macro'])
yknn2 = scores2['test_precision_macro']
print("accuracy for 5-fold cross validation is:")
print(scores2['test_accuracy'])
yknn3 = scores2['test_accuracy']

print("++++++++++++++++++++++++++++++++++++++++++")
##### Gaussian Naive Bayse
scoring = ['precision_macro', 'recall_macro', 'accuracy']
scores3 = cross_validate(clf5, x, y, cv=5, scoring=scoring)
print("recall for 5-fold cross validation in GNB is:")
print(scores3['test_recall_macro'])
ygnb1 = scores3['test_recall_macro']
print("precision for 5-fold cross validation in GNB is:")
print(scores3['test_precision_macro'])
ygnb2 = scores3['test_precision_macro']
print("accuracy for 5-fold cross validationin GNB is:")
print(scores3['test_accuracy'])
ygnb3 = scores3['test_accuracy']

#Plot
#5-Fold Curve
xfold = [1,2,3,4,5]
plt.plot(xfold, yknn1, label='Recall')
plt.plot(xfold, yknn2, label='Precision')
plt.plot(xfold, yknn3, label='Accuracy')
plt.xlabel('5-Fold')
plt.ylabel('Measurment')
plt.title('KNN 5-Fold Curve')
plt.legend(loc='best')
plt.show()

xfold = [1,2,3,4,5]
plt.plot(xfold, ygnb1, label='Recall')
plt.plot(xfold, ygnb2, label='Precision')
plt.plot(xfold, ygnb3, label='Accuracy')
plt.xlabel('5-Fold')
plt.ylabel('Measurment')
plt.title(' Guassian Naive Bayse 5-Fold Curve')
plt.legend(loc='best')
plt.show()

#####ROC Curve
# fpr_knn_lm, tpr_knn_lm = roc_curve(y_test, result2)
# fpr_svc_lm, tpr_svc_lm = roc_curve(y_test, result3)
# fpr_lr_lm, tpr_lr_lm = roc_curve(y_test, result4)
#
# plt.figure(1)
# plt.plot([0, 1], [0, 1], 'k--')
# plt.plot(fpr_knn_lm, tpr_knn_lm, label='KNN')
# plt.plot(fpr_svc_lm, tpr_svc_lm, label='SVC')
# plt.xlabel('False positive rate')
# plt.ylabel('True positive rate')
# plt.title('ROC curve')
# plt.legend(loc='best')
# plt.show()

end_time = datetime.now()
print('Time for running this Code is: {}'.format(end_time - start_time))


import tensorflow as tf
import cv2
import matplotlib.pyplot as plt
import numpy as np

import os
import tarfile
import random
import glob
import shutil

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, Conv2D, Dense, MaxPooling2D, Input, Flatten

# Create Paths for positive, negative and Anchor Images
POSITIVE_PATH = os.path.join("data", "positive")
NEGATIVE_PATH = os.path.join("data", "negative")
ANCHOR_PATH = os.path.join("data", "anchor")

# Create directories for positive and anchor images
if os.path.isdir(POSITIVE_PATH) is False:
    os.makedirs(POSITIVE_PATH)
    os.makedirs(ANCHOR_PATH)

# download link for lfw dataset
# download "All images as gzipped tar file"
# url = http://vis-www.cs.umass.edu/lfw/#download

# Move the lfw.tgz file to the project directory

# extracting the lfw tar gz file
if os.path.isdir('lfw') is False:
    file = tarfile.open('lfw.tgz')
    file.extractall()
    file.close()

# Create negative directory
if os.path.isdir(NEGATIVE_PATH) is False:
    os.makedirs(NEGATIVE_PATH)
    # Moving images from lfw folder to negative folder
    for directory in os.listdir('lfw'):
        for file in os.listdir(os.path.join('lfw', directory)):
            EX_PATH = os.path.join('lfw', directory, file)
            NEW_PATH = os.path.join(NEGATIVE_PATH, file)
            os.replace(EX_PATH, NEW_PATH)

# Accessing the Web cam
cap = cv2.VideoCapture(0)

i = j = 0
while cap.isOpened():
    ret, frame = cap.read()

    # creating 250*250 frames
    frame = frame[120:370, 200:450, :]

    # Press a to collect anchor images
    # Collect anchors -- make sure to collect atleast 200 anchor images
    if cv2.waitKey(1) & 0xFF == 97:
        # create the unique file path
        imgname = os.path.join(ANCHOR_PATH, f'anchor_img{j}.jpg')
        cv2.imwrite(imgname, frame)
        j = j + 1

    # Press p to collect positive images
    # Collect positives -- make sure to collect atleast 200 positive images
    if cv2.waitKey(1) & 0xFF == 112:
        # create the unique file path
        imgname = os.path.join(POSITIVE_PATH, f'positive_img{i}.jpg')
        cv2.imwrite(imgname, frame)
        i = i + 1

    cv2.imshow('Image Collection', frame)

    if cv2.waitKey(1) & 0xFF == 27:
        break

cap.release()
cv2.destroyAllWindows()

# Taking 200 images from each directory
anchor = tf.data.Dataset.list_files(ANCHOR_PATH + '\*.jpg').take(200)
positive = tf.data.Dataset.list_files(POSITIVE_PATH + '\*.jpg').take(200)
negative = tf.data.Dataset.list_files(NEGATIVE_PATH + '\*.jpg').take(200)

# Preprocessing the data
def preprocess(file_path):
    # Read the images
    read_img = tf.io.read_file(file_path)
    # Convert images to tensors
    img = tf.io.decode_jpeg(read_img)
    # Resizing the images to 105 * 105
    img = tf.image.resize(img, (105, 105))
    # Normalizing the images
    img = img / 255.0
    return img

# Creating training data
positives = tf.data.Dataset.zip((anchor, positive, tf.data.Dataset.from_tensor_slices(tf.ones(len(anchor)))))
negatives = tf.data.Dataset.zip((anchor, negative, tf.data.Dataset.from_tensor_slices(tf.zeros(len(anchor)))))
data = positives.concatenate(negatives)

# Preprocessing training data
def preprocess_data(input_img, val_img, label):
    return (preprocess(input_img), preprocess(val_img), label)

# Mapping the data
data = data.map(preprocess_data)
# Shuffling the data
data = data.shuffle(buffer_size=1024)
# Creating the batch of 16
data = data.batch(16)

# Create the model
def create_model():
    inp = Input(shape=(105, 105, 3), name='input_image')
    conv1 = Conv2D(64, (10, 10), activation='relu')(inp)
    max1 = MaxPooling2D((2, 2), padding='same')(conv1)
    conv2 = Conv2D(128, (7, 7), activation='relu')(max1)
    max2 = MaxPooling2D((2, 2), padding='same')(conv2)
    conv3 = Conv2D(128, (4, 4), activation='relu')(max2)
    max3 = MaxPooling2D((2, 2), padding='same')(conv3)
    conv4 = Conv2D(256, (4, 4), activation='relu')(max3)
    flat1 = Flatten()(conv4)
    dense1 = Dense(4096, activation='sigmoid')(flat1)

    return Model(inputs=[inp], outputs=[dense1], name='model')

model = create_model()

# creating L1Dist class to calculate the difference value between two images
class L1Distance(Layer):
    def __init__(self, **kwargs):
        super().__init__()

    def call(self, input_value, valid_value):
        return tf.math.abs(input_value - valid_value)

def create_siamese_model():
    input_image = Input(name='input_image', shape=(105, 105, 3))
    validation_image = Input(name='validation_image', shape=(105, 105, 3))

    siamese_layer = L1Distance()
    distances = siamese_layer(model(input_image), model(validation_image))

    classifier = Dense(1, activation='sigmoid')(distances)

    return Model(inputs=[input_image, validation_image], outputs=classifier, name='siameseNetwork')

siamese_model = create_siamese_model()

binary_loss = tf.losses.BinaryCrossentropy()
optimizer = tf.optimizers.Adam(1e-4)

# Training the model
def train(data, epochs):
    for epoch in range(1, epochs+1):
        print(f'Epoch {epoch}/{epochs}')
        progbar = tf.keras.utils.Progbar(len(data))

        for idx, batch in enumerate(data):
            with tf.GradientTape() as tape:
                x = batch[:2]
                y = batch[2]

                y_pred = siamese_model(x, training=True)
                loss = binary_loss(y, y_pred)
            progbar.update(idx + 1)

            grad = tape.gradient(loss, siamese_model.trainable_variables)
            optimizer.apply_gradients(zip(grad, siamese_model.trainable_variables))

if os.path.isfile('siamesemodel.h5') is False:
    train(data, epochs=50)
    siamese_model.save('siamesemodel.h5')

final_model = tf.keras.models.load_model('siamesemodel.h5', custom_objects={'L1Distance':L1Distance, 'BinaryCrossentropy':tf.losses.BinaryCrossentropy()})

# Creating the path for validation and input images
VAL_PATH = os.path.join('app_data', 'verification_image')
INP_PATH = os.path.join('app_data', 'input_image')

# Creating validation and input directories
if os.path.isdir(VAL_PATH) is False:
    os.makedirs(VAL_PATH)
    os.makedirs(INP_PATH)
    # copying 50 random images from anchor folder to validation folder
    for c in random.sample(glob.glob(os.path.join(ANCHOR_PATH, 'anchor_img*')), 50):
        shutil.copy(c, VAL_PATH)

# Verification function
def verify(model, detection_threshold, verification_threshold):
    results = []
    for image in os.listdir(os.path.join(VAL_PATH)):
        input_img = preprocess(os.path.join(INP_PATH, 'input_image.jpg'))
        validation_img = preprocess(os.path.join('app_data', 'verification_image', image))

        result = model.predict(list(np.expand_dims([input_img, validation_img], axis=1)))
        results.append(result)

    detection = np.sum(np.array(results) > detection_threshold)
    verification = detection / len(os.listdir(os.path.join('app_data', 'verification_image')))
    verified = verification > verification_threshold

    return verified

cap = cv2.VideoCapture(0)
while cap.isOpened():
    ret, frame = cap.read()

    frame = frame[120:370, 200:450, :]

    cv2.imshow('Verification', frame)

    # Press V to verify
    if cv2.waitKey(10) & 0xFF == 118:
        cv2.imwrite(os.path.join('app_data', 'input_image', 'input_image.jpg'), frame)

        verified = verify(final_model, 0.5, 0.5)
        print(verified)

    # Press escape key to escape
    if cv2.waitKey(10) & 0xFF == 27:
        break

cap.release()
cv2.destroyAllWindows()


from django import forms
from .models import Task

class Input_Form(forms.ModelForm):
    class Meta:
        model = Task
        fields = [
            'title',
            'description',
            ]

        labels = {
            'title' : 'Task Name',
            'description' : 'Task Description'
        }

        input_attrs = {
            'type' : 'text',
            'placeholder' : 'Nama task',
        }

        input_attrs2 = {
            'type' : 'text',
            'placeholder' : 'Deskripsi task',
        }

        widgets = {
            'title' : forms.TextInput(attrs=input_attrs),
            'description' : forms.TextInput(attrs=input_attrs2),
        }


import time, random
import numpy as np
from absl import app, flags, logging
from absl.flags import FLAGS
import cv2
import matplotlib.pyplot as plt
import tensorflow as tf
from yolov3_tf2.models import (
    YoloV3, YoloV3Tiny
)
from yolov3_tf2.dataset import transform_images
from yolov3_tf2.utils import draw_outputs, convert_boxes

from deep_sort import preprocessing
from deep_sort import nn_matching
from deep_sort.detection import Detection
from deep_sort.tracker import Tracker
from tools1 import generate_detections as gdet
from PIL import Image
from collections import deque
import velocity
import csv
out_csv = "position.csv"

flags.DEFINE_string('classes', './data/labels/coco.names', 'path to classes file')
flags.DEFINE_string('weights', './weights/yolov3.tf',
                    'path to weights file')
flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny')
flags.DEFINE_integer('size', 516, 'resize images to')
flags.DEFINE_string('video', './data/video/test.mp4',
                    'path to video file or number for webcam)')
flags.DEFINE_string('output', None, 'path to output video')
flags.DEFINE_string('output_format', 'XVID', 'codec used in VideoWriter when saving video to file')
flags.DEFINE_integer('num_classes', 80, 'number of classes in the model')
pts = [deque(maxlen=3) for _ in range(9999)]
count = 1
Real_rate = 1   ## If the tracking accuracy is not 100%, the script counts some vehicles as double.
                ## Therefore, the number of vehicles is greater than the actual number.
                ## This parameter compensates for the effect of tracking accuracy on counting accuracy.
                ## If tracking accuracy is 90%, set this value to 0.9.
def initialize_csv():
    with open(out_csv, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["frame", "vehicle_id", "pixel_x", "pixel_y"])

def main(_argv):
    initialize_csv()
    # Definition of the parameters
    max_cosine_distance = 0.5
    nn_budget = None
    nms_max_overlap = 1.0

    #initialize deep sort
    model_filename = 'model_data/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename, batch_size=1)
    metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
    tracker = Tracker(metric)

    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    if len(physical_devices) > 0:
        tf.config.experimental.set_memory_growth(physical_devices[0], True)

    if FLAGS.tiny:
        yolo = YoloV3Tiny(classes=FLAGS.num_classes)
    else:
        yolo = YoloV3(classes=FLAGS.num_classes)

    yolo.load_weights(FLAGS.weights)
    logging.info('weights loaded')

    class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
    logging.info('classes loaded')

    try:
        vid = cv2.VideoCapture(int(FLAGS.video))
    except:
        vid = cv2.VideoCapture(FLAGS.video)

    out = None

    if FLAGS.output:
        # by default VideoCapture returns float instead of int
        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vid.get(cv2.CAP_PROP_FPS))
        codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
        # out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
        out = cv2.VideoWriter(FLAGS.output, codec, fps, (1024, 768))
        list_file = open('detection.txt', 'w')
        frame_index = -1
    velocity_tracker = velocity.VelocityTracker()
    number_of_frame_in_video = 0
    fps = 0.0
    T_count = 0
    go_count = 0
    come_count = 0
    tails = []
    while True:
        _, img = vid.read()

        if img is None:
            logging.warning("Empty Frame")
            time.sleep(0.1)
            count+=1
            if count < 3:
                continue
            else:
                break
        number_of_frame_in_video += 1

        img = cv2.resize(img, (1024, 768))

        img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_in = tf.expand_dims(img_in, 0)
        img_in = transform_images(img_in, FLAGS.size)

        t1 = time.time()
        boxes, scores, classes, nums = yolo.predict(img_in)
        classes = classes[0]
        names = []
        for i in range(len(classes)):
            names.append(class_names[int(classes[i])])
        names = np.array(names)
        converted_boxes = convert_boxes(img, boxes[0])
        features = encoder(img, converted_boxes)
        detections = [Detection(bbox, score, class_name, feature) for bbox, score, class_name, feature in zip(converted_boxes, scores[0], names, features)]

        #initialize color map
        cmap = plt.get_cmap('tab20b')
        colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

        # run non-maxima suppresion
        boxs = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        classes = np.array([d.class_name for d in detections])
        indices = preprocessing.non_max_suppression(boxs, classes, nms_max_overlap, scores)
        detections = [detections[i] for i in indices]

        # Call the tracker
        tracker.predict()
        tracker.update(detections)

        for track in tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            bbox = track.to_tlbr()
            class_name = track.get_class()
            color = colors[int(track.track_id) % len(colors)]
            color = [i * 255 for i in color]
            # cv2.rectangle(img, (int(bbox[0]), int(bbox[1]-30)), (int(bbox[0])+(len(class_name)+len(str(track.track_id)))*17, int(bbox[1])), color, -1)
            center = (int(((bbox[0])+(bbox[2]))/2),int(((bbox[1])+(bbox[3]))/2))
            pts[int(track.track_id)].append(center[1])
            if len(pts[int(track.track_id)]) > 2 and abs(pts[int(track.track_id)][0] - pts[int(track.track_id)][2]) > 1:
                cv2.putText(img, class_name,(int(bbox[0]), int(bbox[1]-10)),0, 0.6, color,2)
                cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 1)
                cv2.circle(img, (int(((bbox[0])+(bbox[2]))/2),int(((bbox[1])+(bbox[3]))/2)), 2, (0, 0xFF, 0), thickness=2)
            ## Speed estimation and check entering status for counting.
            entered = velocity_tracker.entering_check(fps, number_of_frame_in_video, track.track_id, center, img, bbox)
            if entered == 1:
                if center[0] > 500:
                    go_count = go_count + 1
                else:
                    come_count = come_count + 1

            ## Write position per vehicle and frame
            pos_vicl_fram = [str(number_of_frame_in_video), str(track.track_id), str(center[0]), str(center[1])]
            with open(out_csv, 'a', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(pos_vicl_fram)

        T_count = come_count * Real_rate + go_count * Real_rate
        # print count per frame on console.
        print("Frame: ", number_of_frame_in_video, "  ,  Count: ", nums)
        cv2.putText(img, "Frame: {:.2f}".format(number_of_frame_in_video), (0, 60), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
        cv2.putText(img, "Vehicles/Frame: {:.2f}".format(nums[0]), (0, 90), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
        # print count on screen
        cv2.putText(img, "Total Vehicles: {:.2f}".format(int(T_count)), (500, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (5, 5, 255), 2)
        cv2.putText(img, "Coming Vehicles: {:.2f}".format(int(come_count * Real_rate)), (500, 60), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (5, 5, 255), 2)
        cv2.putText(img, "Going Vehicles: {:.2f}".format(int(go_count * Real_rate)), (500, 90), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (5, 5, 255), 2)
        # print fps on screen
        fps  = ( fps + (1./(time.time()-t1)) ) / 2
        cv2.putText(img, "FPS: {:.2f}".format(fps), (0, 30),
                          cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
        cv2.imshow('output', img)
        if FLAGS.output:
            out.write(img)
            frame_index = frame_index + 1
            list_file.write(str(frame_index)+' ')
            if len(converted_boxes) != 0:
                for i in range(0,len(converted_boxes)):
                    list_file.write(str(converted_boxes[i][0]) + ' '+str(converted_boxes[i][1]) + ' '+str(converted_boxes[i][2]) + ' '+str(converted_boxes[i][3]) + ' ')
            list_file.write('\n')

        # press q to quit
        if cv2.waitKey(1) == ord('q'):
            break
    vid.release()
    if FLAGS.output:
        out.release()
        list_file.close()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass

# python object_tracker.py --video ./data/video/test.mp4 --output ./data/video/results.avi


from setuptools import setup

PACKAGE_NAME = 'rulefinder'

REQUIREMENTS = []
with open('requirements.txt') as f:
    for line in f:
        line, _, _ = line.partition('#')
        line = line.strip()
        REQUIREMENTS.append(line)

setup(
   name=PACKAGE_NAME,
   version='0.1.0',
   description='A project for WSI course at MiNI faculty at WUT',
   packages=[PACKAGE_NAME],
   install_requires= REQUIREMENTS
)


#!/usr/bin/env python3

from pydub import AudioSegment

song = AudioSegment.from_mp3("hyperreal.mp3")

# pydub does things in milliseconds
ten_seconds = 10 * 1000

first_10_seconds = song[ten_seconds:]
last_5_seconds = song[:-5000]

# boost volume by 6dB
beginning = first_10_seconds + 6

# reduce volume by 3dB
end = last_5_seconds - 3

without_the_middle = beginning + end

# song is not modified
backwards = song.reverse()

# 1.5 second crossfade
with_style = beginning.append(end, crossfade=1500)

# repeat the clip twice
do_it_over = with_style * 2

# 2 sec fade in, 3 sec fade out
awesome = do_it_over.fade_in(2000).fade_out(3000)

beat = (song[8800:9300].fade_out(13000).reverse() + 5) * 10
beat.export("mashup.mp3", format="mp3", bitrate="192k")


import mne
from mne.time_frequency import psd_array_multitaper

from itertools import combinations
from scipy.stats import ttest_rel
import scipy
import pandas as pd
import numpy as np
from tqdm.notebook import tqdm
from scipy.interpolate import interp1d

from sklearn.metrics import accuracy_score, f1_score, precision_score
from sklearn.metrics import recall_score

from functools import partial
from pathlib import Path

trim_mean = partial(scipy.stats.trim_mean, proportiontocut=0.1)

def get_spindles_annotations(raw, l_freq=10, h_freq=15, win_width=0.75,
                             ch_name=None, low_th=1.2, high_th=3.5,
                             min_duration=0.5, max_duration=10.0,
                             robust=False):
    resamp_freq = 50

    # 10-15 Hz bandpass
    if ch_name is None:
        ch_name = [ch_name for ch_name in raw.ch_names
                   if "eeg" in ch_name.lower() or "channel" in ch_name.lower()][0]
    bp_data = raw.copy().pick(ch_name).load_data().filter(l_freq, h_freq)
    bp_data = bp_data.resample(resamp_freq)
    times = bp_data.times
    bp_data = bp_data.get_data().squeeze()

    series = pd.Series(bp_data**2, index=times)

    # Rolling windows for MS computation
    rolling = series.rolling(int(win_width*resamp_freq), 1, True)
    if robust:
        res = np.sqrt(rolling.agg(trim_mean))**3  # cubed RMS
    else:
        res = np.sqrt(rolling.mean())**3  # cubed RMS

    # Rolling window for computing the average MS signal amplitude
    page_width = 10*60
    rolling = series.rolling(int(page_width*resamp_freq), 1, True)
    if robust:
        res_page = np.sqrt(rolling.agg(trim_mean))**3  # cubed RMS
    else:
        res_page = np.sqrt(rolling.mean())**3  # cubed RMS

    # Thresholding
    df = pd.DataFrame({"high": res > high_th*res_page,
                       "low": res > low_th*res_page})

    # Making windows based on the low threshold
    mask = df.low.diff().fillna(False)

    # Force transitions if starts or ends with True values
    mask.iloc[0] = df.low.values[0]
    mask.iloc[-1] = df.low.values[-1]

    tmp = df.index[mask]
    starts = tmp[::2]
    stops = tmp[1::2]

    # Selecting only the windows with a peak above the high threshold
    tmp = np.array([(start, stop) for start, stop in zip(starts, stops)
                    if np.sum(df.loc[start:stop, "high"])])
    if not len(tmp):
        return mne.Annotations([], [], [], orig_time=raw.info["meas_date"])

    # Checking window duration based on the low threshold
    durations = np.diff(tmp).squeeze()
    valid_windows = tmp[(durations >= min_duration) &
                        (durations <= max_duration), :]
    if not len(valid_windows):
        return mne.Annotations([], [], [], orig_time=raw.info["meas_date"])

    assert np.all(np.diff(valid_windows).squeeze() >= min_duration)

    return mne.Annotations(valid_windows[:, 0],
                           np.diff(valid_windows).squeeze(),
                           "spindle", orig_time=raw.info["meas_date"])

def get_manual_annotations(file_path, raw):

    spin_day1 = pd.read_excel(file_path)
    spin_day1 = spin_day1.drop(0, axis=0)
    spin_day1 = spin_day1[['Start Time', 'End Time']]

    spin_day1.columns = ["start", "stop"]

    spin_day1["meas_date"] = raw.info["meas_date"]
    spin_day1.start = spin_day1.start.astype(spin_day1.meas_date.dtype)
    spin_day1.stop = spin_day1.stop.astype(spin_day1.meas_date.dtype)

    starts = (spin_day1.start-spin_day1.meas_date)
    starts = starts.astype('timedelta64[ms]')/1000
    stops = (spin_day1.stop-spin_day1.meas_date)
    stops = stops.astype('timedelta64[ms]')/1000

    return mne.Annotations(starts, stops-starts, "spindle_manual",
                           orig_time=raw.info["meas_date"])

def fix_edf(stages_df):
    vals, counts = np.unique(stages_df.start, return_counts=True)
    if len(np.unique(counts)) == 1 and np.unique(counts)[0] == 1:
        return stages_df  # nothing to fix
    for val in vals[np.where(counts != 6)]:
        stages_df = stages_df[stages_df.start != val]
    for i in range(1, 6):
        assert np.all(stages_df.iloc[::6]["start"].values ==
                      stages_df.iloc[i::6]["start"].values)
        stages_df.iloc[i::6].loc[:, "start"] += 10*i
    return stages_df

def get_sleep_stages_df(file_path, raw, stage_duration=10.0):
    if file_path.suffix == ".csv":
        stages = pd.read_csv(file_path)
    else:
        stages = pd.read_excel(file_path)

    stages = stages.drop(index=0)[["Rodent Sleep" if "Rodent Sleep" in stages else "Label",
                                   "Start Time" if "Start Time" in stages else "Time Stamp"]]

    stages.columns = ["label", "start"]
    starts = stages.start.astype('datetime64[ns, UTC]')
    stages.start = (starts - raw.info["meas_date"]).astype('timedelta64[s]')
    if file_path.suffix == ".csv":
        stages = fix_edf(stages)
    return stages

def get_sleep_stages(file_path, raw, stage_duration=10.0):
    stages = get_sleep_stages_df(file_path, raw, stage_duration)

    return np.sum([mne.Annotations(stages.start[stages.label == stage],
                                   stage_duration,
                                   f"stage_{stage}",
                                   orig_time=raw.info["meas_date"])
                   for stage in stages.label.unique()])

def filter_annotations(annotations, tmin, tmax):
    # keep annotation which middle point is within tmin and tmax, without
    # croping them if they party overlap outside of this range
    mid_points = [annot["onset"] + annot["duration"]/2
                  for annot in annotations]
    dat = np.array([(annot["onset"], annot["duration"], annot["description"])
                    for annot, mid_point in zip(annotations, mid_points)
                    if tmin < mid_point < tmax]).T
    if len(dat):
        onset, duration, description = dat
        return mne.Annotations(onset, duration, description,
                               orig_time=annotations.orig_time)
    return mne.Annotations([], [], [], orig_time=annotations.orig_time)

def filter_annots_per_stage(annot_to_filter, annot_stages, stage_in="stage_S"):
    if not len(annot_to_filter):
        return annot_to_filter
    annots = {"onset": [annot["onset"] for annot in annot_stages],
              "duration": [annot["duration"] for annot in annot_stages],
              "description": [annot["description"] for annot in annot_stages]}
    annots = pd.DataFrame(annots)

    cond = annots.description == stage_in
    discontinuities = np.where(np.diff(annots[cond].onset) != 10)[0]+1
    discontinuities = np.concatenate([[0], discontinuities, [np.sum(cond)]])
    block_lenghts = np.diff(discontinuities)*10
    block_starts = annots[cond].onset.values[discontinuities[:-1]]

    lst_annots = [filter_annotations(annot_to_filter, start, start+duration)
                  for start, duration in zip(block_starts, block_lenghts)]
    annotations = np.sum(lst_annots)
    if not isinstance(annotations, mne.Annotations):
        # Happens when there is not annotations
        return lst_annots[0]
    return annotations

def file_list_to_dict(files):
    keys = [(path.parent.name, path.name.split(".")[-2].split("_")[-1])
            for path in files]
    return dict(zip(keys, files))

def get_annotated_raw(edf_file, spindle_file, sleep_file,
                      ch_name=None):
    raw = mne.io.read_raw_edf(edf_file)

    spin_annotations = get_spindles_annotations(raw, ch_name=ch_name)
    annotations_manual = get_manual_annotations(spindle_file, raw)
    annotations_stages = get_sleep_stages(sleep_file, raw)

    spin_annotations = filter_annots_per_stage(spin_annotations,
                                               annotations_stages)

    print([raw.annotations, spin_annotations,
           annotations_manual, annotations_stages])
    raw.set_annotations(raw.annotations + spin_annotations +
                        annotations_manual + annotations_stages)
    return raw

def get_sample_df(raw, sleep_file):
    sample_df = pd.DataFrame({"time": raw.times[::10]})
    sample_df["spindle"] = False
    sample_df["spindle_manual"] = False
    sample_df = sample_df.set_index("time")

    annot_df = raw.annotations.to_data_frame()
    starts = annot_df.onset.astype('datetime64[ns, UTC]')
    annot_df["start"] = (starts-raw.info["meas_date"])
    annot_df["start"] = annot_df["start"].astype('timedelta64[ms]')/1000
    annot_df["stop"] = annot_df["start"] + annot_df.duration

    for _, row in annot_df.iterrows():
        sample_df.loc[row.start:row.stop, row.description] = True

    stages = pd.read_excel(sleep_file)
    stages = stages.drop(index=0)[["Label", "Start Time"]]
    stages.columns = ["label", "start"]
    starts = stages.start.astype('datetime64[ns, UTC]')
    stages.start = (starts-raw.info["meas_date"]).astype('timedelta64[s]')
    stage_duration = 10.0

    sample_df["nrem"] = False
    sample_df["rem"] = False
    for _, row in stages.iterrows():
        if row.label == "S":
            sample_df.loc[row.start:row.start+stage_duration, "nrem"] = True
        if row.label == "P":
            sample_df.loc[row.start:row.start+stage_duration, "rem"] = True

    return sample_df

def get_performances(sample_df):
    return {"accuracy": accuracy_score(sample_df.spindle_manual,
                                       sample_df.spindle),
            "f1": f1_score(sample_df.spindle_manual,
                           sample_df.spindle),
            "precision": precision_score(sample_df.spindle_manual,
                                         sample_df.spindle),
            "recall": recall_score(sample_df.spindle_manual,
                                   sample_df.spindle)}

def get_spin_features(raw, event_name, l_freq=10,
                      h_freq=15, ch_name="EEG 1"):
    bp_data = raw.copy().pick(ch_name).load_data().filter(l_freq, h_freq)

    target_freqs = np.arange(10, 15.1, 0.1)
    p2p = []
    powers = []
    for annot in tqdm(raw.annotations):
        if annot["description"] == event_name:
            dat = bp_data.get_data(tmin=annot["onset"],
                                   tmax=annot["onset"]+annot["duration"])
            p2p.append(dat.max()-dat.min())
            power, freqs = psd_array_multitaper(dat, raw.info["sfreq"],
                                                fmin=5, fmax=20, verbose=False)
            try:
                f = interp1d(freqs, power, kind='cubic')
                powers.append(f(target_freqs).squeeze())
            except ValueError:
                powers.append(target_freqs*np.nan)

    spin_feat = pd.DataFrame([(annot["onset"], annot["duration"])
                              for annot in raw.annotations
                              if annot["description"] == event_name],
                             columns=["onset", "duration"])

    # The two systems recorded in different units
    if np.mean(p2p) > .005:
        spin_feat["p2p_amp"] = np.array(p2p)*1e3  # In uV
    elif np.mean(p2p) < .001:
        spin_feat["p2p_amp"] = np.array(p2p)*1e6  # In uV
    elif len(p2p) == 0:
        spin_feat["p2p_amp"] = []
    else:
        raise ValueError(f"Check units. Mean p2p == {np.mean(p2p)}; p2p={p2p}")

    for freq, power_values in zip(target_freqs, np.array(powers).T):
        spin_feat[str(np.round(freq, 1))] = power_values

    spin_feat["event"] = event_name

    return spin_feat

def get_spin_features_events(raw, ch_name, event_names):
    return pd.concat([get_spin_features(raw, event_name, ch_name=ch_name)
                      for event_name in event_names])

def get_plotting_data(raw, ch_name, nb_hours=4, bins_per_hours=5,
                      event_kinds=("spindle", "spindle_manual")):
    feat = get_spin_features_events(raw, ch_name, event_kinds)

    bins = np.linspace(0, nb_hours*60*60, nb_hours*bins_per_hours+1)
    bin_labels = (bins[1:] + bins[:-1])/2/60/60
    feat["bins"] = pd.cut(feat["onset"], bins=bins,
                          labels=bin_labels, right=False)

    target_freqs = feat.loc[:, "10.0":"15.0"].columns
    groups = feat.drop(columns=target_freqs)\
                 .melt(id_vars=["onset", "event", "bins"])\
                 .groupby(["event", "bins", "variable"])
    data = groups.mean().drop(columns="onset")
    data_count = groups.count().drop(columns="onset").reset_index()
    data_count = data_count[data_count.variable == "duration"]
    data_count["variable"] = "count"

    sleep_durations = {}
    for tmin, tmax, bin in zip(bins[:-1], bins[1:], bin_labels):
        loop_annots = raw.annotations.copy().crop(tmin=tmin, tmax=tmax,
                                                  use_orig_time=False)
        sleep_durations[bin] = np.sum([annot["duration"]
                                       for annot in loop_annots
                                       if annot["description"] == "stage_S"])
    sleep_durations = pd.Series(sleep_durations).reset_index()
    sleep_durations.columns = ["bins", "sleep_duration"]

    sleep_durations.bins = pd.Categorical(sleep_durations.bins.astype(float))

    data_density = data_count.copy().merge(sleep_durations, on="bins")
    data_density["value"] /= data_density.sleep_duration/60
    data_density = data_density.drop(columns="sleep_duration")
    data_density["variable"] = "density"

    data = pd.concat([data,
                      data_count.set_index(["event", "bins", "variable"]),
                      data_density.set_index(["event", "bins", "variable"])])
    data = data.sort_values(['event', 'variable', 'bins'])
    return data.reset_index(), feat.drop(columns=["bins"]).reset_index()

def get_sleep_duration_df(sleep_files):
    sleep_duration = {key: np.sum(pd.read_excel(path).iloc[1:, 0] == "S")/6 for key, path in sleep_files.items()}
    sleep_duration = pd.Series(sleep_duration).reset_index()
    sleep_duration.columns = ["animal", "condition", "nrem_min"]
    #sleep_duration.loc[sleep_duration.condition == "Vehicle 1", "condition"] = "Vehicle"
    #sleep_duration.loc[sleep_duration.condition == "Vehicle 2", "condition"] = "Vehicle"
    sleep_duration.loc[sleep_duration.condition == "Intermittent 1", "condition"] = "Intermittent"
    sleep_duration.loc[sleep_duration.condition == "Intermittent 2", "condition"] = "Intermittent"
    sleep_duration = sleep_duration[sleep_duration.condition != "Intermittent"]
    return sleep_duration.set_index(["animal", "condition"])

def get_pvalues(df, sleep_duration, event_kinds=("spindle",)):
    res_inds = get_stats_ind(df, sleep_duration, event_kinds)

    # Averaging results across Vehicle 1 and Vehicle 2
    res_inds.loc[res_inds.condition == "Vehicle 1", "condition"] = "Vehicle"
    res_inds.loc[res_inds.condition == "Vehicle 2", "condition"] = "Vehicle"
    res_inds = res_inds.groupby(["feature", "condition", "scoring", "animal"]).mean().reset_index()

    p_val_res = []
    for (scoring, feature), df in res_inds.groupby(['scoring', 'feature']):
        for x, y in list(combinations(df.condition.unique(), 2)):
            res = {"scoring": scoring, "feature": feature}
            tmp = df.set_index("animal")  # Aligned by the index when concatenating
            tmp = pd.concat([tmp[tmp.condition == x].rename(columns={"value": x})[x],
                             tmp[tmp.condition == y].rename(columns={"value": y})[y]],
                            axis=1).dropna()
            res["t-value"], res["p-value"] = ttest_rel(tmp[x], tmp[y])
            res["N"] = len(tmp[x])
            res["condition1"] = x
            res["condition2"] = y
            p_val_res.append(res)

    return pd.DataFrame(p_val_res).sort_values(["scoring", "feature"])

def get_stats_ind(df, sleep_duration, event_kinds=("spindle", "spindle_manual")):
    all_res = []
    for scoring in event_kinds:
        for feature in ["duration", "p2p_amp"]:
            tmp = df.groupby(["event", "animal", "condition"]).mean()[[feature]].reset_index()
            tmp = tmp.pivot_table(index=["event", "animal"], columns="condition")[feature]
            if feature == "p2p_amp":
                tmp.loc[tmp[df.condition.unique()[0]] > 5000, df.condition.unique()] /= 1000

            res = tmp.loc[scoring]
            res["scoring"] = scoring
            res["feature"] = feature
            all_res.append(res)

        feature = "duration"
        tmp = df.groupby(["event", "animal", "condition"]).count()[[feature]].reset_index("event")
        tmp = tmp.merge(sleep_duration, left_index=True, right_index=True)
        tmp[feature] /= tmp.nrem_min
        tmp = tmp.drop(columns="nrem_min").rename(columns={"duration": "density"}).reset_index()
        feature = "density"
        tmp = tmp.pivot_table(index=["event", "animal"], columns="condition")[feature]

        res = tmp.loc[scoring]
        res["scoring"] = scoring
        res["feature"] = feature
        all_res.append(res)

    stats_df = pd.concat(all_res).reset_index().melt(id_vars=["animal", "scoring", "feature"])
    return stats_df.sort_values(["scoring", "feature", "condition"])

def get_stats(df, sleep_duration, event_kinds=("spindle", "spindle_manual")):
    res_inds = get_stats_ind(df, sleep_duration, event_kinds)

    # Averaging results across Vehicle 1 and Vehicle 2
    res_inds.loc[res_inds.condition == "Vehicle 1", "condition"] = "Vehicle"
    res_inds.loc[res_inds.condition == "Vehicle 2", "condition"] = "Vehicle"
    res_inds = res_inds.groupby(["feature", "condition", "scoring", "animal"]).mean().reset_index()

    return pd.concat([res_inds.groupby(["feature", "condition", "scoring"]).mean().rename(columns={"value": "mean"})["mean"],
                      res_inds.groupby(["feature", "condition", "scoring"]).std().rename(columns={"value": "std"})["std"],
                      res_inds.groupby(["feature", "condition", "scoring"]).count().rename(columns={"value": "count"})["count"]],
                     axis=1)

def get_spin_feat_df(df_perf, file_pattern="{animal}_{condition}_spin_feat.csv"):
    dfs = []
    for animal, condition in df_perf[["animal", "condition"]].values:
        path = Path(".") / file_pattern.format(animal=animal,
                                               condition=condition)
        if not path.exists():
            continue
        df = pd.read_csv(path, index_col=0)
        df['animal'], df['condition'] = path.name.split("_")[:2]
        dfs.append(df)

    df = pd.concat(dfs)
    #df.loc[df.condition == "Vehicle 1", "condition"] = "Vehicle"
    #df.loc[df.condition == "Vehicle 2", "condition"] = "Vehicle"
    df.loc[df.condition == "Intermittent 1", "condition"] = "Intermittent"
    df.loc[df.condition == "Intermittent 2", "condition"] = "Intermittent"
    return df[df.condition != "Intermittent"]

def get_perf_df(study_no, edf_files, sleep_files,
                spindle_files, recompute=False,
                file_pattern="detection_performance_study{study_no}.csv"):
    file_name = Path(file_pattern.format(study_no=study_no))
    if recompute or not file_name.exists():
        results = []
        for key in tqdm(list(edf_files)):
            raw = mne.io.read_raw_edf(edf_files[key])
            ch_names = raw.pick("eeg").ch_names
            ch_names = [ch_name for ch_name in ch_names
                        if "eeg" in ch_name.lower() or "channel" in ch_name.lower()]
            for ch_name in ch_names:
                raw = get_annotated_raw(edf_files[key], spindle_files[key],
                                        sleep_files[key], ch_name)
                sample_df = get_sample_df(raw, sleep_files[key])
                res = get_performances(sample_df)
                res["animal"] = key[0]
                res["condition"] = key[1]
                res["ch_name"] = ch_name
                results.append(res)

        df_perf = pd.DataFrame(results)
        df_perf = df_perf.sort_values('f1', ascending=False).drop_duplicates(['animal','condition'])
        df_perf = df_perf.reset_index(drop=True).sort_values("animal")
        df_perf.to_csv(file_name)
        return df_perf.drop(columns="ch_name")

    df_perf = pd.read_csv(file_name, index_col=0).reset_index(drop=True)
    return df_perf.sort_values("animal")

def get_whole_raw(edf_file, condition, animal, start_times, ch_name):
    raw = mne.io.read_raw_edf(edf_file)
    tmin = (start_times[(animal, condition)] - start_times[(animal, "whole")]).days*24*60*60
    tmin += (start_times[(animal, condition)] - start_times[(animal, "whole")]).seconds
    tmin +=  (start_times[(animal, condition)] - start_times[(animal, "whole")]).microseconds/10e6
    tmax = tmin + 12*60*60
    cropped_raw = raw.copy().crop(tmin, tmax)
    cropped_raw.set_meas_date(start_times[(animal, condition)])
    cropped_raw._cropped_samp = 0
    return cropped_raw

def get_start_times(whole_edf_files, edf_files):
    start_times = {}
    for animal in tqdm(whole_edf_files):
        for condition in ["baseline", "kynurenine"]:
            raw = mne.io.read_raw_edf(edf_files[(animal, condition)])
            start_times[(animal, condition)] = raw.info["meas_date"]

        raw = mne.io.read_raw_edf(whole_edf_files[animal])
        start_times[(animal, "whole")] = raw.info["meas_date"]

    return start_times


from bs4 import BeautifulSoup
from urllib.parse import quote_plus
import requests

APPS = []

def download(link):
	res = requests.get(link + '/download?from=details', headers={
			'User-Agent':'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/601.7.5 (KHTML, like Gecko) Version/9.1.2 Safari/601.7.5'
		}).text
	soup = BeautifulSoup(res, "html.parser").find('a', {'id':'download_link'})
	if soup['href']:
		r = requests.get(soup['href'], stream=True, headers={
			'User-Agent':'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/601.7.5 (KHTML, like Gecko) Version/9.1.2 Safari/601.7.5'
		})
		with open(link.split('/')[-1] + '.apk', 'wb') as file:
			for chunk in r.iter_content(chunk_size=1024):
				if chunk:
					file.write(chunk)

def search(query):
	res = requests.get('https://apkpure.com/search?q={}&region='.format(quote_plus(query)), headers={
			'User-Agent':'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/601.7.5 (KHTML, like Gecko) Version/9.1.2 Safari/601.7.5'
		}).text
	soup = BeautifulSoup(res, "html.parser")
	for i in soup.find('div', {'id':'search-res'}).findAll('dl', {'class':'search-dl'}):
		app = i.find('p', {'class':'search-title'}).find('a')
		APPS.append((app.text,
					i.findAll('p')[1].find('a').text,
					'https://apkpure.com' + app['href']))


import pandas as pd
from elasticsearch import AsyncElasticsearch
from dotenv import load_dotenv
import os

load_dotenv()
name = os.getenv("NAME")
es = AsyncElasticsearch(hosts=f"http://{name}:9200")
mapping = {"properties": {"id": {"type": "integer"}, "text": {"type": "text"}}}

class Elastic:
    async def search(query):
        resp = await es.search(
            index="documents",
            body={"query": {"match": {"text": query}}},
            size=20,
        )
        arr = [i["_source"]["id"] for i in resp["hits"]["hits"]]
        return arr

    async def delete(id):
        try:
            await es.delete_by_query(
                index="documents", body={"query": {"match": {"id": id}}}
            )
            return
        except:
            return

    async def fill():
        if not await es.indices.exists(index="documents"):
            await es.indices.create(
                index="documents", body={"mappings": mapping}
            )
            df = pd.read_csv("posts.csv")
            df = df.reset_index()
            df = df.rename(columns={"index": "id"})
            documents = df.to_dict(orient="records")
            for document in documents:
                await es.index(
                    index="documents",
                    body={"id": document["id"] + 1, "text": document["text"]},
                )