from django.db import models
from django.utils.translation import gettext as _
from django import forms
from django.forms.widgets import TextInput
from django.db.models import Q
from django.db.models import Sum, DurationField, ExpressionWrapper, F, IntegerField, Value

import math
import datetime
from utils import time

# from .widgets import FileSignField

from courses.models import Courses
from users.models import User
from countries.models import Country, City
from centers.models import Center
from institutions.models import Institution
from digisafe.storage import FileInstCertStorage

### Manage files ###
import os
from django.core.exceptions import ValidationError
def validate_file_size(value):
    filesize= value.size
    if filesize > 2097152: #2MB
        raise ValidationError(_("The maximum file size that can be uploaded is {0}".format("2MB")))
    else:
        return value

def validate_file_extension(value):
    ext = os.path.splitext(value.name)[1]  # [0] returns path+filename
    valid_extensions = ['.pdf']
    if not ext.lower() in valid_extensions:
        raise ValidationError(_("Unsupported file extension. Valid extension are '{0}'".format("pdf")))

def file_path_name(instance, file_name):
    file_root, file_ext = os.path.splitext(file_name)
    return "prot_{0}_{1}_{2}".format(instance.protocol.id, instance.doc_type, file_ext)

from django.core.files.storage import Storage, FileSystemStorage
from PyPDF2 import PdfFileWriter, PdfFileReader
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4, landscape
import io, shutil

class ProtocolFileSystemStorage(FileSystemStorage):
    def _save(self, name, content):
        # print("ProtocolFileSystemStorage._save name", name)
        file_root, file_ext = os.path.splitext(name)
        name = self.get_alternative_name(file_root, file_ext)
        # print("ProtocolFileSystemStorage._save name", name)
        return super()._save(name, content)

    def delete(self, name):
        # print("ProtocolFileSystemStorage.delete name", name)
        super().delete(name)
        signedName = self.checkSigned(name)
        if signedName:
            super().delete(signedName)

    def signFile(self, name, fileObj=None):
        file_root, file_ext = os.path.splitext(name)
        signedFilePath = file_root+"_signed"+file_ext
        # if self.exists(signedFilePath):
            # print("exist", signedFilePath)
            # master = signedFilePath
        # else:
            # master = name
        master = name
        self.drawSigns(master, signedFilePath, fileObj)

    def checkSigned(self, name):
        # ritorna il nome del file se esiste, altrimenti False
        file_root, file_ext = os.path.splitext(name)
        signedFile = file_root+"_signed"+file_ext
        if self.exists(signedFile):
            return signedFile
        return False

    def getSignedFileName(self, name):
     # name:  deve essere sempre il file originale
        if not name or type(name) != type(""):
            return False
        # print("ProtocolFileSystemStorage.getSigned", name)
        file_root, file_ext = os.path.splitext(name)
        signedFile = file_root+"_signed"+file_ext
        if self.exists(signedFile):
            return signedFile
        return name

    def drawSigns(self, signFile, destFile, fileObj):
        # signFile: il file che deve essere firmato
        #           se esiste una versione firmata prende quest'ultima
        #           altrimenti usa la versione originale da firmare
        # destFile: il file con le firme messe

        signFile_tmp = self.path(signFile+".tmp")
        destFile = self.path(destFile)
        # todo: bisogna migliorare il controllo del processo per evitare due richieste contemporanee
        if self.exists(signFile_tmp):
            print("Processo bloccato. File tmp esistente. Attendere la fine del processo in corso")
            # os.remove(self.path(signFile_tmp))
            return False
        shutil.copyfile(self.path(signFile), self.path(signFile_tmp))
        # read your existing PDF
        sign_file = open(signFile_tmp, "rb")
        existing_pdf = PdfFileReader(sign_file)
        output = PdfFileWriter()
        pages = existing_pdf.getNumPages()
        signs = Signs.objects.filter(file=fileObj).order_by("uiid")

        for pag in range(0, pages):
            page = existing_pdf.getPage(pag)

            packet = io.BytesIO()
            can = canvas.Canvas(packet, pagesize=A4)

            if signs.filter(page=pag+1):
                for sign in signs.filter(page=pag+1):
                    mask = [215,255,215,255,215,255] # elimina il bianco
                    # print(sign.user.profile.sign.path)
                    can.drawImage(sign.user.profile.sign.path, sign.pdf_x-50, sign.pdf_y-25, 100, 50, mask=mask, anchor="c", preserveAspectRatio =True)
                    can.drawCentredString (sign.pdf_x, sign.pdf_y, "%s %s"%(sign.user.first_name,sign.user.last_name) )
                    can.drawCentredString (sign.pdf_x, sign.pdf_y-14, "%s"%(sign.timestamp.strftime("%d-%m-%Y %H:%M")) )

            else:
                can.drawString(0,0, "")
            can.save()
            #move to the beginning of the StringIO buffer
            packet.seek(0)
            # create a new PDF with Reportlab
            new_pdf = PdfFileReader(packet)
            page.mergePage(new_pdf.getPage(0))

            output.addPage(page)
        # finally, write "output" to a real file
        outputStream = open(destFile, "wb")
        output.write(outputStream)
        outputStream.close()
        sign_file.close()
        # Rimuove la copia del file
        os.remove(self.path(signFile_tmp))

    def url(self, name):
        # Ritorna il file firmato se esiste
        # print("ProtocolFileSystemStorage.url")
        return super().url(self.getSignedFileName(name))

fs = ProtocolFileSystemStorage(location='files/', base_url="/document")
### END Manage files ###

class Protocol(models.Model):
    datetime   = models.DateTimeField(auto_now_add=True)
    owner      = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    # editable=False,
                    )
    course     = models.ForeignKey(
                    Courses,
                    on_delete=models.CASCADE,
                )

    TYPE_CHOICES = [
        ("new", _('New')),
        ("update", _('Update')),
        ]
    type        = models.CharField(
                    max_length=25,
                    choices=TYPE_CHOICES,
                    default="new",
                )

    learners_request = models.IntegerField(default=0)
    STATUS_CHOICES = [
        ("m", 'Modifica'),
        ("r", 'Richiesta'),
        ("c", 'Caricato'),
        ("a", 'Autorizzato'),
        ("n", 'Negato'),
        ("t", 'Terminato'),
        ("h", 'Chiuso'),
        ]
    status      = models.CharField(
                    max_length=1,
                    choices=STATUS_CHOICES,
                    default="m",
                )

    center      = models.ForeignKey(
                    Center,
                    on_delete=models.CASCADE,
                    blank=True,
                    null=True,
                )
    institution =  models.ForeignKey(
                    Institution,
                    on_delete=models.CASCADE,
                    blank=True,
                    null=True,
                )
    warning     = models.BooleanField(default=False, editable=False)
    def __str__(self):
        return "Prot. {num}".format(num=self.pk)

    def title(self):
        return self.course.feature.title

    def code(self):
        return self.course.code

    def getTrainers(self):
        return [x.trainer for x in self.session_set.all()]

    def getLearners(self):
        return [x.user for x in self.learners_set.all()]

    def checkHoursPracticeRequestMin(self):
        return self.getHoursPractice() - self.getHoursPracticeRequestMin()

    def getHoursPracticeRequestMin(self):
        h,m = time.timedelta_to_hm(getattr(self.course, self.type).practice)
        return h

    def getHoursTheoryRequest(self):
        h,m = time.timedelta_to_hm(getattr(self.course, self.type).theory)
        return h

    def getHoursTheory(self):
        ore_teoria_inserite = self.session_set.filter(subject_type="t").aggregate(
                    total_time=Sum(
                        ExpressionWrapper(F('end_time') - F('start_time'), output_field=DurationField())
                    )
                )['total_time']
        h,m = time.timedelta_to_hm(ore_teoria_inserite)
        return h

    def getDifferenceHoursTheory(self):
        # print("Protocol.getDifferenceHoursTheory", int(self.getHoursTheoryRequest())-int(self.getHoursTheory()))
        return int(self.getHoursTheoryRequest())-int(self.getHoursTheory())

    def checkHoursTheory(self):
        ore_teoria_inserite = self.getHoursTheory()
        ore_teoria = self.getHoursTheoryRequest()
        # print("checkHoursTheory", ore_teoria_inserite >= ore_teoria)
        return ore_teoria_inserite >= ore_teoria

    def getLearnersRequest(self):
        # return self.learners_set.count()
        return self.learners_request

    def checkNumTrainer(self):
        maxLearnersByTrainer = self.getPracticeLearnersPerTainer()
        if maxLearnersByTrainer <= 0: return True # nessun rapporto docenti/discenti
        trainers = self.session_set.filter(subject_type="p").count()
        learners = self.getLearnersRequest()

        min_trainer = int(math.ceil(learners/maxLearnersByTrainer))
        # print("Learners.checkNumTrainer", "pk:", self.id,
                                          # "trainers:", trainers,
                                          # "learners:", learners,
                                          # "maxLearnersByTrainer:", maxLearnersByTrainer)
        # print("Learners.checkNumTrainer min_trainer", min_trainer,
                                        # "Num trainer ok?", int(trainers) >= int(min_trainer) )
        return (int(trainers) >= int(min_trainer))

    def getMaxLearnersTheory(self):
        typeCourse = self.type
        return getattr(self.course, typeCourse).max_learners_theory

    def getMaxLearnersPractice(self):
        return self.getPracticeLearnersPerTainer()

    def getPracticeLearnersPerTainer(self):
        typeCourse = self.type
        return getattr(self.course, typeCourse).max_learners_practice

    def getAuthorizationMessage(self):
        if self.course.need_institution:
            return  _("Training project authorized by the trainer {0} with protocol {1} of {2}.".format(self.institution, "xxx", "YYY"))
        return  _("Activities planned by the employer and provided through a qualified person pursuant to the D.M. March 6, 2013.")

    def getRequalifictionMessage(self):
        return  _("Qualification update within {0} years.".format(self.course.feature.years))

    def getQualification(self):
        """ Data conseguimento qualifica """
        return self.session_set.latest('date').date

    def getExpiration(self):
        """ Data scadenza qualifica """
        years = self.course.feature.years
        qual  = self.getQualification() #session_set.latest('date').date
        return qual.replace(year = qual.year + years)

    def expired(self):
        # print(type(self.getExpiration()))
        now = datetime.date.today()
        return self.getExpiration() < now

    def checkPractice(self):
        h = self.getHoursPracticeRequestMin()
        return h>0

    def checkHoursPractice(self):
        if not self.checkNumTrainer(): return False
        hoursInsertPractice = self.getHoursPractice()
        h = self.getHoursPracticeRequest()
        # print("Protocol.checkHoursPractice", h, "learners:",self.learners_set.count(), "Ore/doc", self.getPracticeLearnersPerTainer())
        return (hoursInsertPractice >= h)

    def getDifferenceHoursPractice(self):
        return int(self.getHoursPractice()) - int(self.getHoursPracticeRequest())

    def getHoursPracticeRequest(self):
        learners = self.getLearnersRequest()
        hoursPractice,m = time.timedelta_to_hm(getattr(self.course, self.type).practice)
        if self.getPracticeLearnersPerTainer() <= 0: return 0 # nessun discente previsto
        h = hoursPractice * math.ceil(learners / self.getPracticeLearnersPerTainer() )
        return h

    def getHoursPractice(self):
        hours = self.session_set.filter(subject_type="p").aggregate(
                    total_time=Sum(
                        ExpressionWrapper(F('end_time') - F('start_time'), output_field=DurationField()),
                    )
                )['total_time']
        h,m = time.timedelta_to_hm(hours)
        # print("Protocol.getHoursPractice", h, hours)
        return h

    def checkCenter(self):
        return self.center or False

    def checkAll(self):
        # print("Protocol.checkAll",self.checkHoursTheory , self.checkHoursPractice , self.checkNumTrainer)
        return self.checkHoursTheory() and self.checkHoursPractice() and self.checkNumTrainer() and self.checkCenter()

    def checkAllSignedFiles(self):
        # print(self.files_set.all())
        # print([x.signs_set.all().count()>0 for x in self.files_set.all()])
        if False in [x.signs_set.all().count()>0 for x in self.files_set.all()]:
            # print("False")
            return False
        return True

    def checkAllCertificateLoads(self):
        # print("checkAllCertificateLoads", [x.inst_cert.name!="" for x in self.learners_set.filter(passed=True)])
        if False in [x.inst_cert.name!="" for x in self.learners_set.filter(passed=True)]:
            return False
        return True

    def status_close(self):
        return self.status == "h"

    def is_authorized(self):
        if self.status in ['a', 't', 'h']:
            return True
        return False

class Session(models.Model):
    protocol    = models.ForeignKey(
                    Protocol,
                    on_delete=models.CASCADE,
                )
    trainer     = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    null=True,
                    # blank=True
                )
    country     = models.ForeignKey(
                    Country,
                    on_delete=models.CASCADE,
                    null=True,
                    # blank=True
                )
    city        = models.ForeignKey(
                    City,
                    on_delete=models.CASCADE,
                    null=True,
                    # blank=True
                )
    address     = models.CharField(max_length=255, default="",null=True)
    TYPE_CHOICES = [
        ("t", _('Theory')),
        ("p", _('Practice')),
        ]
    subject_type    = models.CharField(
                    max_length=25,
                    choices=TYPE_CHOICES,
                    default="t",
                )
    EXECUTION_CHOICES = [
        ("front", _('In presence')),
        ("online", _('On the web')),
        ]
    execution   =  models.CharField(
                    max_length=25,
                    choices=EXECUTION_CHOICES,
                    default="front",
                )
    date        = models.DateField(null=True)
    start_time  = models.TimeField(null=True)
    end_time    = models.TimeField(null=True)

    # @property
    # def type(self):
        # return self.subject_type

    def __str__(self):
        return _("Session {data} {inizio}/{fine}".format(data=self.date,
                                            inizio=self.start_time,
                                            fine=self.end_time
                                        ))

    def getPlace(self):
        return "{0} {1}".format(self.city, self.address)

    def getTheoryTimes(self, date, start_time, end_time):
        "Sovrapposizione di orari nella stessa data per la teoria"
        # return
        return Session.objects.filter(
                                (Q(start_time__lt=start_time) & Q(end_time__gt=start_time))
                                |
                                (Q(start_time__lt=end_time) & Q(end_time__gt=end_time))
                                |
                                (Q(start_time__gt=start_time) & Q(start_time__lt=end_time))
                                |
                                (Q(end_time__gt=start_time) & Q(end_time__lt=end_time)),
                                subject_type="t", date=date, pk=self.pk,
                                )  #.exclude()

    def getDateTimeRange(self, trainer, date, start_time, end_time):
        "Impegni del Trainer diviso per data ed orario"
        return Session.objects.filter( trainer=trainer,
                                date=date,
                                start_time__lte=start_time,
                                end_time__gte=end_time).exclude(pk=self.pk)

    def isSignedInProtocol(self):
        learners = self.protocol.learners_set.all()
        for learner in learners:
            if learner.user == self.trainer:
                return self.trainer
        return False

    def isBusyInstance(self):
        return self.isBusy(self.trainer)

    def isBusy(self, learner):
        sessions = self.protocol.session_set.all()
        for session in sessions:
            date = session.date
            start_time = session.start_time
            end_time = session.end_time
            if Protocol.objects.filter( learners__user=learner,
                                session__date=date,
                                session__start_time__lte=start_time,
                                session__end_time__gte=end_time
                                ).exclude(pk=self.protocol.pk):
                return True
        return False

    def isPreDate(self, date):
        # print("Session isPreDate")
        if self.protocol.course.need_institution:
            pre_days = self.protocol.institution.pre_days
            today = datetime.date.today()
            end_date = today + datetime.timedelta(days=pre_days)

            # print(date, type(date))
            # print(today, type(today))
            # print(end_date, type(end_date))
            # print(date >= end_date)
            return date >= end_date
        else:
            return True

def file_path_name_learner(instance, file_name):
    file_root, file_ext = os.path.splitext(file_name)
    return "prot_{0}_{1}_{2}".format(instance.protocol.id, "cert_inst", file_ext)
class Learners(models.Model):
    protocol    = models.ForeignKey(
                    Protocol,
                    on_delete=models.CASCADE,
                )
    user        = models.ForeignKey(
                User,
                on_delete=models.CASCADE,
                null=True, blank=True
            )
    passed      = models.BooleanField(default=False)
    inst_cert   = models.FileField(upload_to=file_path_name_learner,
                                   storage=FileInstCertStorage(location='files/', base_url="/document"),
                                   blank=True,
                                   validators=[validate_file_size, validate_file_extension],
                                   verbose_name=_("Certificate Institution")
                                  )

    class Meta:
        ordering = ['user__last_name']
        unique_together = ('protocol', 'user',)

    def __str__(self):
        return _("User: {nome} {cognome} ({user}) ".format(user=self.user.username,
                                            nome=self.user.first_name,
                                            cognome=self.user.last_name
                                        ))

    def isSignedInProtocol(self):
        sessions = self.protocol.session_set.all()
        for session in sessions:
            if self.user == session.trainer:
                return self.user
        return False

    def isBusyInstance(self):
        return self.isBusy(self.user)

    def isBusy(self, learner):
        sessions = self.protocol.session_set.all()
        for session in sessions:
            date = session.date
            start_time = session.start_time
            end_time = session.end_time
            # print(date, start_time, end_time )
            if Protocol.objects.filter( learners__user=learner,
                                session__date=date,
                                session__start_time__lte=start_time,
                                session__end_time__gte=end_time
                                ).exclude(pk=self.protocol.pk):
                return True
        return False

    # todo: cancella il file se cancello il record ( ma non lo avevo fatto ??)

class Files(models.Model):
    protocol    = models.ForeignKey(
                    Protocol,
                    on_delete=models.CASCADE,
                )
    datetime    = models.DateTimeField(auto_now_add=True)
    owner       = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    editable=False,
                    )
    TYPE_CHOICES = [
        ("r", _('Attendance register')),
        ("v", _('Exam reporter')),
        ("p", _('Privacy')),
        ("o", _('Other')),
        ]
    doc_type    = models.CharField(
                    max_length=2,
                    choices=TYPE_CHOICES,
                    default="r",
                    # required=True,
                )
    file        = models.FileField(upload_to=file_path_name,
                                   storage=fs, blank=False,
                                   validators=[validate_file_size, validate_file_extension],
                                   # help_text="file help text"
                                  )
    class Meta:
        ordering = ['id']
        unique_together = ('protocol', 'doc_type',)

    def delete(self, *args, **kwargs):
        self.file.delete()
        super().delete(*args, **kwargs)

    def __str__(self):
        # print("Files.__str__", type(self.file))
        return "Files"+self.get_doc_type_display()

    def check_signed(self):
        if self.signs_set.all():
            return True
        return False

    def check_file(self):
        # print(self.file.storage.url(self.file.name))
        # print(dir(self.file.storage))
        if self.file.name:
            return True
        else:
            return False

    def getSignedFile(self):
        return self.file.storage.getSignedFileName(self.file.name)

    def get_signed_or_master(self):
        if self.file.storage.getSignedFileName(self.file.name):
            return self.file.storage.getSignedFileName(self.file.name)
        else:
            return self.file.name

    def signFile(self, u=None, data=None):
        # print("Files.signFile")
        # print(self.file.name)
        # Genera il file se esiste almeno una firma associata
        if self.signs_set.all():
            self.file.storage.signFile(self.file.name, self)
        else:
            #altrimenti cancella il file "xxxx_signed.pdf" se esiste
            if self.file.storage.checkSigned(self.file.name):
                self.file.storage.delete(self.file.storage.checkSigned(self.file.name))

    def _openSignFile(self):
        from django.utils.safestring import SafeString
        from django.urls import reverse_lazy

        html = "<a class='digi-btn digi-bg-primary' href='{0}'>{1}</a>".format(
                    reverse_lazy("admin:sign-file", kwargs={'pk':self.protocol.id, 'file_id':self.id}),
                    _("Sign"))
        return SafeString(html)

    def html(self):
        return self._openSignFile()

# Authorizations by Institutions
class Authorizations(models.Model):
    protocol    = models.ForeignKey(
                    Protocol,
                    on_delete=models.CASCADE,
                )
    datetime    = models.DateTimeField(auto_now_add=True)
    owner       = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    editable=False,
                    )
    TYPE_CHOICES = [
        ("a", _('Authorization')),
        # ("o", _('Other informations')),
        ]
    doc_type    = models.CharField(
                    max_length=2,
                    choices=TYPE_CHOICES,
                    default="a",
                    # required=True,
                )
    auth_prot   = models.CharField(max_length=20, default="")
    file        = models.FileField(upload_to=file_path_name,
                                   storage=fs, blank=False,
                                   validators=[validate_file_size, validate_file_extension],
                                  )
    class Meta:
        ordering = ['id']
        unique_together = ('protocol', 'doc_type',)

    def delete(self, *args, **kwargs):
        self.file.delete()
        super().delete(*args, **kwargs)

    def __str__(self):
        # print("Files.__str__", type(self.file))
        return "Authorizations"

class Action(models.Model):
    protocol    = models.ForeignKey(
                    Protocol,
                    on_delete=models.CASCADE,
                )
    owner       = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    editable=False,
                    )
    datetime    = models.DateTimeField(auto_now_add=True)
    text        = models.CharField(max_length=255, default="")
    info        = models.TextField(default="")

import uuid
class Signs(models.Model):
    uiid = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    timestamp = models.DateTimeField(auto_now_add=True)
    file = models.ForeignKey(
                    Files,
                    on_delete=models.CASCADE,
                )
    user = models.ForeignKey(
                    User,
                    on_delete=models.CASCADE,
                    editable=False,
                )

    pdf_x     =  models.IntegerField(default=0)
    pdf_y     =  models.IntegerField(default=0)
    html_x    =  models.IntegerField(default=0)
    html_y    =  models.IntegerField(default=0)
    page      =  models.IntegerField(default=0)


"""
ant
***

:Author: tobijjah
:Date: 31.05.19
"""
from abc import ABCMeta
from abc import abstractmethod

from math import sqrt
from numpy import array
from pygame import Surface

from ant.settings import ALPHA
from ant.settings import ANT_WITHOUT_NUTRIENT_COLOR, ANT_WITH_NUTRIENT_COLOR
from ant.settings import BETA
from ant.settings import GAMMA
from ant.settings import GLOBAL_RNG
from ant.settings import Q

class BaseAnt(metaclass=ABCMeta):

    INSTANCES = 0

    def __init__(
            self,
            cell,
            background,
            rect,
            gamma=GAMMA,
            beta=BETA,
            alpha=ALPHA,
            q=Q
    ):

        self._rng = GLOBAL_RNG
        self._name = __class__.INSTANCES
        __class__.INSTANCES += 1

        self.pos = cell.pos
        self.background = background

        width, height = 0.2*rect.width, 0.2*rect.height

        self.surface = Surface((width, height))
        self.rect = background.get_rect(
            top=self._rng.uniform(rect.top, rect.bottom - height),
            left=self._rng.uniform(rect.left, rect.right - width),
            width=width,
            height=height
        )
        self.x_off, self.y_off = self.rect.left - rect.left, self.rect.top - rect.top

        # algorithm parameter
        self._gamma = gamma
        self._alpha = alpha
        self._beta = beta
        self._q = q

        self._mandible = 0
        self._current_cell = cell
        self._movement_cell = None

        self._visited = list()  # contains visited cells
        self._allowed = list()  # allowed movement cells

        self._trail_length = 0

    @abstractmethod
    def collide(self, cells):
        pass

    @abstractmethod
    def move(self, cells):
        pass

    def mandible_full(self):
        return self._mandible > 0

    def draw(self):
        if self.mandible_full():
            self.surface.fill(ANT_WITH_NUTRIENT_COLOR)

        else:
            self.surface.fill(ANT_WITHOUT_NUTRIENT_COLOR)

        self._update_rect()
        self.background.blit(self.surface, self.rect)

    def _update_rect(self):
        dx = (self._current_cell.rect.left + self.x_off) - self.rect.left
        dy = (self._current_cell.rect.top + self.y_off) - self.rect.top

        self.rect.move_ip(dx, dy)

    def _allowed_cells(self, cells):
        # filter for obstacles
        self._allowed = list(set(cells) - set(self._visited))
        self._allowed = list(filter(lambda cell: not cell.has_obstacle(), self._allowed))

    def _probabilities(self, target):
        probabilities = array([
            cell.pheromone.intensity ** self._alpha * (1 / __class__.euclidean(cell, target)) ** self._beta
            if cell.has_pheromone() else
            self._gamma ** self._alpha * (1 / __class__.euclidean(cell, target)) ** self._beta
            for cell in self._allowed
        ])

        return probabilities / probabilities.sum()

    @staticmethod
    def euclidean(origin, target):
        return sqrt(
            ((origin.pos.x - target.pos.x)**2 + (origin.pos.y - target.pos.y)**2) + 0.0000001
        )

    def __hash__(self):
        return hash((self.__class__.__name__, self._name))

    def __str__(self):
        return '{}{} POS{} -> NUTRIENT {}u'.format(self.__class__.__name__, self._name,
                                                   str(self.pos), self._mandible)

    def __repr__(self):
        msg = '<{}(pos={}, gamma={}, alpha={}, beta={}, q={}) at {}>'.format(
            self.__class__.__name__, str(self.pos), self._gamma, self._alpha,
            self._beta, self._q, hex(id(self))
        )
        return msg

class SimpleAnt(BaseAnt):
    def __init__(self, cell, rect, background, **kwargs):
        super().__init__(cell, rect, background, **kwargs)

    def collide(self, cells):
        pass

    def move(self, cells):
        if self.mandible_full():
            self._update_pheromone()
            self._movement_cell = self._visited.pop()

            if self._movement_cell.has_hole() and self._movement_cell.hole.home(self):
                self._movement_cell.hole.nutrients += self._mandible
                self._mandible -= 1
                self._trail_length = 0

        else:
            self._visited.append(self._current_cell)
            self._allowed_cells(cells)
            self._select()

            if self._movement_cell.has_nutrient() and not self._movement_cell.nutrient.empty():
                self._mandible += self._movement_cell.nutrient.nutrient_unit

        self._current_cell = self._movement_cell
        self.pos = self._current_cell.pos

    def _select(self):
        if self._allowed:
            self._movement_cell = self._rng.choice(
                self._allowed,
                size=1,
                p=self._probabilities(self._current_cell)
            )[0]

            self._trail_length += __class__.euclidean(self._current_cell, self._movement_cell)

        else:
            self._movement_cell = self._visited[0]
            self._visited = []
            self._trail_length = 0

    def _update_pheromone(self):
        if not self._current_cell.has_pheromone():
            self._current_cell.spawn_pheromone()

        self._current_cell.pheromone.update(self._trail_length)

class Ant:
    # we shoud limit the stored path traveled

    INSTANCES = 0

    def __init__(
            self,
            cell,
            rect,
            surface,
            gamma=GAMMA,
            beta=BETA,
            alpha=ALPHA,
            q=Q
    ):

        self._name = __class__.INSTANCES
        __class__.INSTANCES += 1
        self._rng = GLOBAL_RNG

        self.pos = cell.pos
        self.rect = Rect(rect.centerx, rect.centery, rect.width/2, rect.height/2)
        self.surface = surface

        # algorithm parameter
        self._gamma = gamma
        self._alpha = alpha
        self._beta = beta
        self._q = q

        self._mandible = 0
        self._current_cell = cell
        self._movement_cell = None

        self._idx = -1  # traverse index
        self._dead_end = False  # Ant can get stuck within its own path, if this happen we traverse till Ant is unstuck

        self._visited = []  # contains visited cells
        self._allowed = []  # allowed movement cells

        self.radius = 8
        self.goal = None
        self.length = 0

    def draw(self):
        square(self.surface, ANT_WITHOUT_NUTRIENT_COLOR, self.rect)

    def mandible_full(self):
        return self._mandible > 0

    def collide(self, cells):
        # select goal by nearest neighbor
        if not self.goal:
            for cell in cells:
                if cell.has_nutrient() and self._circle_collide(cell):
                    self.goal = cell
                    print('bam')

    def move(self, cells):
        if self.mandible_full():
            self.update()
            self.back()

        else:
            self._allowed_cells(cells)

            if self._dead_end:
                self._traverse()

            else:
                self._visited.append(self._current_cell)
                self._select()

        self._move()

    def update(self):
        if self._current_cell.has_pheromone():
            self._current_cell.pheromone.update(self.length)

        else:
            self._current_cell.pheromone = Pheromone(self._current_cell.rect, self._current_cell.surface)
            self._current_cell.pheromone.update(self.length)

    def back(self):
        self._movement_cell = self._visited.pop()

    def _move(self):
        previous = self._current_cell
        self._current_cell = self._movement_cell
        self.pos = self._current_cell.pos
        self.rect = Rect(self._current_cell.rect.centerx, self._current_cell.rect.centery,
                         self._current_cell.rect.width/2, self._current_cell.rect.height/2)

        if not self.mandible_full() and self._current_cell.has_nutrient() and not self._current_cell.nutrient.empty():
            self._mandible += self._current_cell.nutrient.nutrient_unit

        if self.mandible_full() and self._current_cell.has_hole():
            self._current_cell.hole.nutrients += self._mandible
            self._mandible -= 1
            self._length = 0
            self.update()

        else:
            self.length += __class__.euclidean(previous, self._current_cell)

    def _allowed_cells(self, cells):
        # filter for obstacles
        self._allowed = list(set(cells) - set(self._visited))

    def _traverse(self):
        self._movement_cell = self._visited[0]
        self._visited = []
        self.length = 0
        # if self._allowed:
        #     self._visited = self._visited[:self._idx + 1]
        #     self._dead_end = False
        #     self._select()
        #
        # else:
        #     self._idx -= 1
        #     self._movement_cell = self._visited[self._idx]

    def _select(self):
        if self._allowed and self.goal:
            self._movement_cell = self._rng.choice(
                self._allowed,
                size=1,
                p=self._probabilities(self.goal)
            )[0]

        elif self._allowed:
            self._movement_cell = self._rng.choice(
                self._allowed,
                size=1,
                p=self._probabilities(self._current_cell)
            )[0]

        else:
            self._idx = -1
            #self._dead_end = True
            self._traverse()

    def _probabilities(self, target):
        probabilities = array([
            cell.pheromone.intensity ** self._alpha * (1 / __class__.euclidean(target, cell)) ** self._beta
            if cell.has_pheromone() else
            self._gamma ** self._alpha * (1 / __class__.euclidean(target, cell)) ** self._beta
            for cell in self._allowed
        ])

        return probabilities / probabilities.sum()

    @staticmethod
    def euclidean(origin, target):
        return sqrt(
            ((origin.pos.x - target.pos.x)**2 + (origin.pos.y - target.pos.y)**2) + 0.0000001
        )

    def _circle_collide(self, cell):
        return (cell.pos.x - self.pos.x)**2 + (cell.pos.y - self.pos.y)**2 <= self.radius**2

    def __hash__(self):
        return hash((self.__class__.__name__, self._name))

    def __str__(self):
        return '{} {} at position {} with {} unit nutrient'.format(self.__class__.__name__,
                                                                   self._name, str(self.pos),
                                                                   self._mandible)

    def __repr__(self):
        msg = '<{}(pos={}, gamma={}, alpha={}, beta={}, q={}) at {}>'.format(self.__class__.__name__, str(self.pos),
                                                                             self._gamma, self._alpha,
                                                                             self._beta, self._q,
                                                                             hex(id(self)))
        return msg


import requests
from bs4 import BeautifulSoup
from pokemon import DualTypePokemon, SingleTypePokemon

def fetch_data_from_webdb():
    dex_page = requests.get('https://pokemondb.net/pokedex/all')
    with open('data.html', 'w', encoding='utf-8') as file:
        file.write(dex_page.text)

def fetch_data_from_localdb():
    dex_page = ''
    with open('data.html', 'r', encoding='utf-8') as file:
        for i in file:
            dex_page += i.strip()
    return dex_page
    pass

if __name__ == '__main__':
    dex_page = fetch_data_from_localdb()
    with open('data.html', 'r', encoding='utf-8') as file:
        for i in file:
            dex_page += i.strip()
    dex_soup = BeautifulSoup(dex_page, 'html.parser')
    dex_table = dex_soup.find('table', attrs={'id': 'pokedex'})
    dex_tbody = dex_table.find('tbody')
    pokemons = dex_tbody.find_all('tr')
    pokemon_list = []
    for i in pokemons:
        attrs = i.find_all('td')
        national_id = attrs[0].find('span', attrs={'class': 'infocard-cell-data'}).text
        national_id = int(national_id)
        name_tag = attrs[1]
        name = name_tag.find('a')
        morph = name_tag.find('small')
        if morph is not None:
            name = morph.text
        else:
            name = name.text
        stats = []
        for j in attrs[4:]:
            stats.append(int(j.text))
        type_tag = attrs[2]
        type_ = type_tag.find_all('a')
        type1_ = None
        type2_ = None
        if len(type_) == 2:
            type1_ = type_[0].text
            type2_ = type_[1].text
            pokemon = DualTypePokemon(name, national_id, type1_, type2_)
            pokemon.set_stats(stats)
        else:
            type1_ = type_[0].text
            pokemon = SingleTypePokemon(name, national_id, type1_)
            pokemon.set_stats(stats)
        pokemon_list.append(pokemon)
    print(pokemon_list[3].id_)
    pass


import json
from typing import Dict, List, Optional, Tuple

import requests
from django.db import transaction
from unidecode import unidecode

from .. import models

class ResponseDictEncoder:
    def encode(self, o):
        if isinstance(o, requests.Response):
            is_json = True
            try:
                content = o.json()
            except json.decoder.JSONDecodeError:
                content = o.content.decode('utf-8')
                is_json = False
            return {
                "url": o.url,
                "status_code": o.status_code,
                "reason": o.reason,
                "headers": dict(o.headers),
                "encoding": o.encoding,
                "is_json": is_json,
                "content": content
            }

class DetachedPublication:
    def __init__(self,
                 publication: models.Publication,
                 author_author_alias_pairs: List[Tuple[models.Author, models.AuthorAlias]],
                 container_container_alias_pair: Tuple[models.Container, models.ContainerAlias],
                 raw: models.Raw,
                 audit_command: models.AuditCommand):

        # these must not be saved in the database
        self.publication = publication
        self.container = container_container_alias_pair[0]
        self.container_alias = container_container_alias_pair[1]
        self.authors_author_alias_pairs = author_author_alias_pairs
        self.raw = raw

        # audit command must already be saved in the database
        self.audit_command = audit_command

    def attach_to(self, publication_id: int, match: Optional[int]):
        with transaction.atomic():
            publication = models.Publication.objects.get(id=publication_id)
            if match is not None:
                self.raw.value["match_ids"] = [match]

            container, container_alias = update_container(
                audit_command=self.audit_command,
                detached_container=self.container,
                detached_container_alias=self.container_alias,
                existing_container=publication.container)

            publication.container = container
            update_publication(publication, self.publication, self.audit_command)

            self.raw.publication = publication
            self.raw.container = container
            self.raw.save()

            for (author, author_alias) in self.authors_author_alias_pairs:
                author_db = models.Author.objects.log_create(audit_command=self.audit_command,
                                                             given_name=author.given_name,
                                                             family_name=author.family_name,
                                                             orcid=author.orcid,
                                                             email=author.email)
                author_alias_db = models.AuthorAlias.objects.log_create(audit_command=self.audit_command,
                                                                        given_name=author_alias.given_name,
                                                                        family_name=author_alias.family_name,
                                                                        author_id=author_db.id)

                models.RawAuthors.objects.log_create(
                    self.audit_command,
                    author_id=author_db.id, raw_id=self.raw.id)

def update_container(existing_container: Optional[models.Container],
                     detached_container: models.Container,
                     detached_container_alias: models.ContainerAlias,
                     audit_command: models.AuditCommand):
    if existing_container:
        if existing_container.issn == '' and detached_container.issn != '':
            existing_container.log_update(audit_command=audit_command, issn=detached_container.issn)
        container = existing_container
    else:
        container = models.Container.objects.log_create(audit_command=audit_command,
                                                        name=detached_container.name,
                                                        issn=detached_container.issn,
                                                        type=detached_container.type)

    container_alias, created = models.ContainerAlias.objects.log_get_or_create(
        audit_command,
        name=detached_container_alias.name,
        container_id=container.id)

    return container, container_alias

def update_publication(
        existing_publication: models.Publication,
        detached_publication: models.Publication,
        audit_command: models.AuditCommand):
    changes = {}
    if existing_publication.date_published_text is None:
        existing_publication.date_published_text = detached_publication.date_published_text
        changes["date_published_text"] = detached_publication.date_published_text
    if not existing_publication.title:
        existing_publication.title = detached_publication.title
        changes["title"] = detached_publication.title
    if not existing_publication.doi:
        existing_publication.doi = detached_publication.doi
        changes["doi"] = detached_publication.doi
    if not existing_publication.abstract:
        existing_publication.abstract = detached_publication.abstract
        changes["abstract"] = detached_publication.abstract

    existing_publication.log_update(audit_command=audit_command, **changes)

def get_message(response_json):
    return response_json["message"]

def make_author_author_alias_pair(publication: models.Publication, item_json: Dict,
                                  create) -> models.Author:
    family = unidecode(item_json.get("family").upper())
    given = unidecode(item_json.get("given").upper())
    orcid = item_json.get("ORCID", "")

    author = models.Author(type=models.Author.INDIVIDUAL,
                           orcid=orcid,
                           given_name=given,
                           family_name=family)
    author_alias = models.AuthorAlias(author=author,
                                      given_name=given,
                                      family_name=family)

    if create:
        author.save()
        author_alias.author = author
        author_alias.save()
    return author, author_alias

def make_author_author_alias_pairs(publication_raw, response_json: Dict, create) -> List[models.Author]:
    authors_json = response_json.get("author", [])
    if not authors_json:
        authors_json = response_json.get("editor", [])
    author_author_alias_pairs = [make_author_author_alias_pair(publication_raw, author_json, create) for author_json in
                                 authors_json]
    return author_author_alias_pairs

def get_year(item_json):
    return item_json \
        .get("issued", {"date-parts": [[None]]}) \
        .get("date-parts", [[None]])[0][0]

def get_title(item_json):
    title = ""
    title_list = item_json.get("title")
    if title_list:
        title = title_list[0]
    return title

def get_doi(item_json):
    return item_json.get("DOI", "")

def get_container_type(item_json):
    return item_json.get("type", "")

def get_container_name(item_json):
    container_title = ""
    container_title_list = item_json.get("container-title")
    if container_title_list:
        container_title = container_title_list[0]
    return container_title

def make_container_container_alias_pair(publication_raw: models.Publication, item_json, create):
    container_name = get_container_name(item_json)
    container_type = get_container_type(item_json)

    container = models.Container(type=container_type, name=container_name)
    container_alias = models.ContainerAlias(name=container_name)

    if create:
        container.save()
        container_alias.container = container
        container_alias.save()
    return container, container_alias


#coding: utf-8
import os
import cv2
import sys
import json
import argparse
import subprocess
import numpy as np
from PIL import Image

from ..utils._colorings import toBLUE, toGREEN, toACCENT
from ..utils.argparse_utils import ListParamProcessorCreate
from ..utils.audio_utils import synthesize_audio
from ..utils.generic_utils import now_str
from ..utils.monitor_utils import ProgressMonitor
from ..utils.pil_utils import draw_text_in_pil

def video_of_lyric(argv=sys.argv[1:]):
    """Create a lyric Video.

    Args:
        json (str)            : Path to parameter json file.
        --ttfontname (str)    : A filename or file-like object containing a TrueType font. If the file is not found in this filename, the loader may also search in other directories, such as the ``fonts/`` directory on Windows or ``/Library/Fonts/`` , ``/System/Library/Fonts/`` and ``~/Library/Fonts/`` on macOS.
        --margin (int)        : The margin size.
        --mode (str)          : Optional mode to use for color values.
        --fontsize (int)      : The font size.
        --fontwidth (int)     : The font width.
        --fontheight (int)    : The font height.
        --img-size (tuple)    : The image size.
        --bgRGB (tuple)       : he color of background image. (RGB)
        --textRGB (tuple)     : The color of text. (RGB)
        --alpha-range (float) : How many seconds to set alpha to 1 (maximum).
        --fps (float)         : The video fps.
        --span (int)          : The span between lyrics.

    .. code-block:: python

        >>> # Create a json.
        >>> import numpy as np
        >>> def func(word, start, end, indent='\\t'*4):
        ...     print(f'{indent}"words": "{word}",')
        ...     print(f'{indent}"seconds": {[round(e,3) for e in np.linspace(start, end, len(word))]},')

    .. code-block:: python

        >>> import json
        >>> from pycharmers.utils import dumps_json
        >>> with open("era-it-doo/era-it.json") as f:
        ...     data = json.load(f)
        >>> print(dumps_json(data))
        {
          "kwargs": {
            "ttfontname": "/Users/iwasakishuto/Library/Fonts/851MkPOP_002.ttf",
            "img_size": [
              360,
              640
            ],
            "fontsize": 38,
            "fontwidth": 30,
            "margin": 10,
            "bgRGB": [
              0,
              0,
              0,
              255
            ],
            "mode": "RGBA",
            "ret_position": "line"
          },
          "texts": [
            [{"words":"１日今日も、","seconds":[2.4, 2.624, 2.848, 3.072, 3.296, 3.52],"x":-1,"y":50},{"words":"こなした労働、","seconds":[4.1, 4.317, 4.533, 4.75, 4.967, 5.183, 5.4],"x":-1,"y":-1},{"words":"辛かったけど","seconds":[5.7, 5.96, 6.22, 6.48, 6.74, 7.0],"x":-1,"y":-1},{"words":"本当がんばったじゃん","seconds":[7.0, 7.222, 7.444, 7.667, 7.889, 8.111, 8.333, 8.556, 8.778, 9.0],"x":-1,"y":-1},{"words":"１日今日も、","seconds":[9.1, 9.38, 9.66, 9.94, 10.22, 10.5],"x":-1,"y":260},{"words":"いちいちどうこう、","seconds":[10.8, 10.962, 11.125, 11.288, 11.45, 11.612, 11.775, 11.938, 12.1],"x":-1,"y":-1},{"words":"いわれたけど","seconds":[12.3, 12.52, 12.74, 12.96, 13.18, 13.4],"x":-1,"y":-1},{"words":"本当にがんばったじゃん","seconds":[13.4, 13.59, 13.78, 13.97, 14.16, 14.35, 14.54, 14.73, 14.92, 15.11, 15.3],"x":-1,"y":-1}]
          ]
        }

    Note:
        When you run from the command line, execute as follows::

        $ video_of_lyric dodo-era-it.json --audio dodo-era-it.mp4[.mp3]

    +--------------------------------------------+
    |                Sample                      |
    +============================================+
    | .. image:: _images/cli.video_of_lyric.gif  |
    +--------------------------------------------+
    """
    parser = argparse.ArgumentParser(prog="video_of_lyric", description="Create a lyric Video.", add_help=True)
    parser.add_argument("json",                type=str, help="Path to parameter json file.")
    parser.add_argument("--ttfontname",  type=str, default=None, help="A filename or file-like object containing a TrueType font. If the file is not found in this filename, the loader may also search in other directories, such as the ``fonts/`` directory on Windows or ``/Library/Fonts/`` , ``/System/Library/Fonts/`` and ``~/Library/Fonts/`` on macOS.")
    parser.add_argument("--margin",      type=int, default=None, help="The margin size.")
    parser.add_argument("--mode",        type=str, default=None, help="Optional mode to use for color values.")
    parser.add_argument("--fontsize",    type=int, default=None, help="The font size.")
    parser.add_argument("--fontwidth",   type=int, default=None, help="The font width.")
    parser.add_argument("--fontheight",  type=int, default=None, help="The font height.")
    parser.add_argument("--img-size",    action=ListParamProcessorCreate(type=int), default=None, help="The image size.")
    parser.add_argument("--bgRGB",       action=ListParamProcessorCreate(type=int), default=None, help="The color of background image. (RGB)")
    parser.add_argument("--textRGB",     action=ListParamProcessorCreate(type=int), default=None, help="The color of text. (RGB)")
    parser.add_argument("--alpha-range", type=float, help="The video length [s].", default=1.)
    parser.add_argument("--fps",         type=float, help="The video fps.", default=30.)
    parser.add_argument("--span",        type=int,   help="The span between lyrics", default=None)
    parser.add_argument("--audio",       type=str,   help="The audio path.", default=None)
    args = parser.parse_args(argv)

    json_path = args.json
    with open(json_path, mode="r") as f:
        data = json.load(f)

    args_kwargs = dict(args._get_kwargs())
    data_kwargs = data.get("kwargs", {})
    get_kwargs = lambda x,default=None:args_kwargs.get(x) or data_kwargs.get(x, default)
    img_size   = tuple(get_kwargs("img_size", default=[360,640])[:2])
    bgRGB      = tuple(get_kwargs("bgRGB", default=(0,0,0,255)))
    textRGB    = tuple(get_kwargs("textRGB", default=(255,255,255,255)))
    mode       = get_kwargs("mode", default="RGBA")
    margin     = get_kwargs("margin", default=10)
    fontsize   = get_kwargs("fontsize", default=40)
    fontwidth  = get_kwargs("fontwidth", default=None)
    fontheight = get_kwargs("fontheight", default=None)
    span       = get_kwargs("span", default=1)
    audio_path = get_kwargs("audio", default=None)
    alpha_range = args.alpha_range
    fps = args.fps
    spf = 1/fps
    kwargs = {
        "ttfontname": get_kwargs("ttfontname", default=""),
        "img_size": img_size,
        "fontsize": fontsize,
        "fontwidth": fontwidth,
        "fontheight": fontheight,
        "margin": margin,
        "bgRGB": bgRGB,
        "mode": mode,
        # "ret_position": "line",
    }
    text_data = data.get("texts", [[]])
    num_texts = len(text_data)
    sec_filter = []
    for i,page in enumerate(text_data):
        start, end = (1e16, 0)
        for line in page:
            secs = line.get("seconds", [])
            start = min(start, *secs)-alpha_range
            end   = max(end, *secs)
        sec_filter.append([i,start,end+alpha_range+span])
    duration = max([e[2] for e in sec_filter]) + alpha_range + 1
    num_frames = int(duration*fps)

    def set_pos(data, name, default):
        p = data.get(name, -1)
        if p==-1: p = default
        return p

    def set_fontcolor(start, sec, textRGB, alpha_range):
        fc = [0]*4
        fc[:3] = textRGB[:3]
        fc[3] = min(255, max(0, int((sec-start)/alpha_range*255)))
        return tuple(fc)

    params = {
        "duration" : duration,
        "alpha range" : alpha_range,
        "FPS": fps,
        "span": span,
        "audio": audio_path,
    }
    params.update(kwargs)
    print("[Parameters]")
    for k,v in params.items():
        print(f"* {toACCENT(k)}: {toGREEN(v)}")

    root, ext = os.path.splitext(json_path)
    video_path = f"{root}_{now_str()}.mp4"
    fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
    out_video = cv2.VideoWriter(video_path, fourcc, fps, img_size)
    monitor = ProgressMonitor(max_iter=num_frames, barname="Editing")
    it = sec = 0
    while True:
        it+=1
        bg = Image.new(mode=mode, size=img_size, color=bgRGB)
        sec += spf
        for i,*_ in filter(lambda x:x[1]<= sec <x[2], sec_filter):
            texts = text_data[i]
            init_x = x = y = margin
            for j,line in enumerate(texts):
                words = line.get("words", "")
                startSecs = line.get("seconds", [])
                x = set_pos(data=line, name="x", default=x)
                y = set_pos(data=line, name="y", default=y)
                for w,start in zip(words, startSecs):
                    fc = set_fontcolor(start, sec, textRGB, alpha_range)
                    bg,(x,_) = draw_text_in_pil(text=w, img=bg, ret_position="word", textRGB=fc, x=x, y=y, **kwargs)
                y += fontsize
                x = init_x
            break
        out_video.write(np.asanyarray(bg.convert("RGB")))
        monitor.report(it, sec=f"{sec:.2f}/{duration}", text=f"{j+1}/{num_texts}")
        if sec>duration:
            break
    out_video.release()
    monitor.remove()
    print(f"{toBLUE(video_path)} (No Sound) is created.")

    if audio_path is not None:
        synthesize_audio(video_path=video_path, audio_path=audio_path)


# -*- coding: utf-8 -*-
import numpy as np
from numpy.linalg import norm
import torch
from torch.utils.data import Dataset

class CQTsDataset(Dataset):
    """CQTs dataset."""

    def __init__(self, n_files, n_triplets=16, bias=True, normalization='max', delta=(16, 1, 96), dim_cqt=(72, 64)):

        self.n_files = n_files
        self.delta = delta
        self.bias = bias
        self.normalization = normalization
        self.n_triplets = n_triplets
        self.dim_cqt = dim_cqt

    def __len__(self):
        return len(self.n_files)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        try:
            cqts = np.load(self.n_files[idx])
        except ValueError:
            print("An error occured with file {}".format(self.n_files[idx]))

        # Normalization
        if self.normalization == 'max':
            cqts = cqts / np.max(cqts + 1e-7)
        elif self.normalization == 'log_max_centered':
            cqts = np.log(cqts + 5e-3)
            cqts = cqts - np.mean(cqts)
            cqts = cqts / np.max(np.abs(cqts))

        cqts_batch = torch.empty(self.n_triplets, 3, self.dim_cqt[0], self.dim_cqt[1])
        for j in range(self.n_triplets):

            a, p, n = self.sample_triplet(cqts)
            cqts_batch[j, 0] = torch.from_numpy(self.append_cqts(cqts, a - 2, a + 2))
            cqts_batch[j, 1] = torch.from_numpy(self.append_cqts(cqts, p - 2, p + 2))
            cqts_batch[j, 2] = torch.from_numpy(self.append_cqts(cqts, n - 2, n + 2))

        return cqts_batch

    def sample_triplet(self, cqts):
        dp, dnmin, dnmax = self.delta
        L = np.shape(cqts)[0]
        a = torch.randint(1, L-1, (1,)).item()
        p_inf = max(a - dp, 0)
        nl_inf = max(a - dnmax, 0)
        nr_inf = min(a + dnmin, L)
        p_sup = min(a + dp, L)
        nl_sup = max(a - dnmin, 0)
        nr_sup = min(a + dnmax, L)

        if self.bias:
            sample_left = self.compare_segments(cqts, a)
            (p_inf, p_sup) = (a + 1, p_sup) if sample_left else (p_inf, a - 1)
        else:
            sample_left = torch.randint(2, (1,)).item()

        p, n = a, a
        while (a == p or a == n):
            p = torch.randint(p_inf, p_sup + 1, (1,)).item()
            n = torch.randint(nl_inf, nl_sup + 1, (1,)).item() if sample_left \
                else torch.randint(nr_inf, nr_sup + 1, (1,)).item()

        return a, p, n

    def append_cqts(self, cqts, low, high):
        N, F, T = np.shape(cqts)
        cqts_app = np.zeros((F, (high - low) * T))
        for i, pos  in enumerate(range(low, high)):
            if 0 <= pos < N:
                cqts_app[:, i*T:(i+1)*T] = cqts[pos]
        return cqts_app

    def compare_segments(self, cqts, a):
        eps = 1e-8
        cqts_left1 = np.log(np.abs(np.fft.fft2(np.log(np.abs(self.append_cqts(cqts, a - 20, a - 12)) + eps))) + eps)
        cqts_left2 = np.log(np.abs(np.fft.fft2(np.log(np.abs(self.append_cqts(cqts, a - 8, a)) + eps))) + eps)
        cqts_right1 = np.log(np.abs(np.fft.fft2(np.log(np.abs(self.append_cqts(cqts, a, a + 8)) + eps))) + eps)
        cqts_right2 = np.log(np.abs(np.fft.fft2(np.log(np.abs(self.append_cqts(cqts, a + 12, a + 20)) + eps))) + eps)

        return norm(cqts_left2 - cqts_left1, 2) > norm(cqts_right2 - cqts_right1, 2)


import inspect
from dataclasses import dataclass

class MagicList(list):
    def __init__(self, cls_type):
        if not inspect.isclass(cls_type):
            raise NotImplementedError('Need to provide type')
        super().__init__()
        self._cls_type = cls_type
        #self._list = list()

    def __getitem__(self, ii):
        if ii < super().__len__():
            return super().__getitem__(ii)
        elif ii == super().__len__():
            new_el = self._cls_type()
            super().append(new_el)
            return super().__getitem__(ii)
        else:
            raise IndexError('List index out of range')

    def __setitem__(self, ii, val):
        if ii < super().__len__():
            super().__setitem__(ii, val)
        elif ii == super().__len__():
            new_el = self._cls_type()
            super().append(new_el)
        else:
            raise IndexError('List index out of range')

    # def __len__(self):
    #     """List length"""
    #     return len(self._list)

    # def __iter__(self):
    #     for elem in self._list:
    #         yield elem

    # def __repr__(self):
    #     return f"{type(self).__name__}({super().__repr__()})"

@dataclass
class Person:
    age: int = 1

# new_element_class = globals()[self._cls_type]
#             new_el = new_element_class()
#             self._list.append(new_el)
#             return new_el


'''
Tools for the computation of
Horton-Strahler stream order.
'''

import matplotlib.pyplot as plt
import numpy as np
import copy

def ReadTree(ccoFile):
    print("Reading", ccoFile)

    with open(ccoFile, 'r') as f:
        # Unused '*Tree' information
        row = f.readline()
        row = f.readline()
        row = f.readline()

        # Vessels data
        f.readline().strip()
        nVessels = int(f.readline())

        treeData = []
        for i in range(nVessels):
            row = (f.readline()).split() # Split all columns in a list
            Id, xProx, xDist, r, q, stage = int(row[0]), row[1:4], row[4:7], float(row[12]), float(row[13]), int(row[-1])
            l = sum([(float(a)-float(b))**2 for a,b in zip(xProx, xDist)])**.5
            treeData.append([Id, r*10, l*10, q, stage]) # Convert to mm

        row = f.readline()
        f.readline().strip()

        treeConnectivity = []
        for i in range(nVessels):
            row = (f.readline()).split()
            if (len(row)==2):
                Id, parentId = int(row[0]), int(row[1])
                treeConnectivity.append([Id, parentId, []])
            else:
                Id, parentId, children = int(row[0]), int(row[1]), [int(i) for i in row[2:]]
                treeConnectivity.append([Id, parentId, children])

        print("\tNumber of vessels in the tree:", nVessels, "\n")

    return treeData, treeConnectivity

def BifurcationOrder(treeData, treeConnectivity, orderMax=1e6):

    treeConnectivity.sort()
    treeData.sort()
    orders = np.zeros((len(treeConnectivity),))+1

    def updateTree(vessel, currentOrder=0):
        # Go up the tree and update all orders
        parentId = int(vessel[1])

        orders[vessel[0]] = max(currentOrder, orders[vessel[0]])
        if parentId:
            parent = treeConnectivity[parentId]
            if len(parent[-1])==2:
                updateTree(parent, min(currentOrder+1,orderMax))
            else:
                updateTree(parent, min(currentOrder, orderMax))
        else:
            orders[parentId] = min(currentOrder, orderMax)
            return

    for v in treeConnectivity:
        if len(v[-1])==0:       # if terminal segment
            updateTree(v, currentOrder=0)

    # Add stream order to vessel data
    treeDataNoRoot = []
    count = 0
    for vessel in treeData:
        Id = int(vessel[0])
        order = orders[Id]
        treeData[Id].append(order)

        if treeData[Id][-2]>=-1:
            Id, radius, length, flow, stage, order = treeData[Id]
            treeDataNoRoot.append([count, radius, length, flow, stage, order])
            count+=1

    return treeDataNoRoot

def StrahlerOrder(treeData, treeConnectivity):

    treeConnectivity.sort()
    currentOrder = 0
    vessel = treeConnectivity[0]
    vesselDict = {vessel[0]: [vessel[1], vessel[2], currentOrder]}
    for vessel in treeConnectivity[1:]:
        if len(vessel)>2:
            IdVessel, IdParent, ChildrenList = vessel
        else:
            IdVessel, IdParent = vessel
            ChildrenList = []
        vesselDict[IdVessel] = [IdParent, ChildrenList, currentOrder]
    prunedTree = copy.deepcopy(vesselDict)

    def CurrentLeaves(prunedTree):
        leavesList = []
        for Id, vessel in prunedTree.items():
            if vessel[-1] == 0 and len(vessel[-2])==0:
                leavesList.append(Id)
        return leavesList

    def UpdateAndPrune(prunedTree, leaves):
        # Assign an order to the current leaves
        copyTree = copy.deepcopy(prunedTree)
        increment= False

        for leaf in leaves:
            Children = treeConnectivity[leaf][-1]
            if len(Children)==2 and copyTree[Children[0]][-1]==currentOrder and copyTree[Children[1]][-1]==currentOrder:
                copyTree[leaf][-1] = currentOrder+1
                increment = True

            else:
                childrenOrders = [0]
                for child in Children:
                    childrenOrders.append(copyTree[child][-1])
                '''
                Two different rules, the max rule is the one from Wikipedia and gives fewer orders.
                Both give the right results for the test tree.
                '''
                copyTree[leaf][-1] = currentOrder # max(1,max(childrenOrders))

        # Prune
        for leaf in leaves:
            parentId = copyTree[leaf][0]
            if parentId>=0:
                copyTree[parentId][1].remove(leaf)
        return copyTree, increment

    currentOrder = 1
    computeForTheRoot = iter([True, False]) # Just used to run one more time once the root is reached
    while prunedTree[0][-1]==0 or next(computeForTheRoot):
        leafList = CurrentLeaves(prunedTree)
        prunedTree, increment = UpdateAndPrune(prunedTree, leafList)
        if increment: currentOrder+=1

        if 0 in leafList:
            break

    # Add stream order to vessel data
    treeDataNoRoot = []
    count = 0
    for vessel in treeData:
        Id = int(vessel[0])
        order = prunedTree[Id][-1]
        treeData[Id].append(order)

        if treeData[Id][-2]>=-1:
            Id, radius, length, flow, stage, order = treeData[Id]
            treeDataNoRoot.append([count, radius, length, flow, stage, order])
            count+=1
    return treeDataNoRoot

def TreeStatistics(orderedTree, treeData):
    return #stats

def PlotTreeStatistics(orderedTree, outputImageFile=None):

    maxOrder = np.array(orderedTree)[:,-1].astype(int).max()
    data = np.array(orderedTree)

    dataRadius = np.zeros((maxOrder, len(orderedTree)))
    dataLength = np.zeros((maxOrder, len(orderedTree)))
    dataVolume = np.zeros((maxOrder, len(orderedTree)))
    dataFlow   = np.zeros((maxOrder, len(orderedTree)))
    dataAspectRatio   = np.zeros((maxOrder, len(orderedTree)))
    orderDistribution = np.zeros((maxOrder,))

    for Id, radius, length, flow, stage, order in orderedTree:
        r,l = radius*1e5, length*1e5 # Convert to microns
        i,j = int(order)-1, int(Id)
        dataRadius[i, j] = radius
        dataLength[i, j] = length
        dataFlow[i,j]    = flow
        dataVolume[i, j] = radius*radius*length*np.pi
        dataAspectRatio[i, j] = length/radius
        orderDistribution[i] += 1

    meanRadius = np.mean(dataRadius, axis = 1, where=dataRadius>0)
    stdRadius  = np.std(dataRadius, axis = 1, where=dataRadius>0)
    meanLength = np.mean(dataLength, axis = 1, where=dataRadius>0)
    stdLength  = np.std(dataLength, axis = 1, where=dataRadius>0)
    meanFlow   = np.mean(dataFlow, axis = 1, where=dataRadius>0)
    stdFlow    = np.std(dataFlow, axis = 1, where=dataRadius>0)
    meanVolume = np.mean(dataVolume, axis = 1, where=dataRadius>0)
    stdVolume  = np.std(dataVolume, axis = 1, where=dataRadius>0)
    meanAspectRatio = np.mean(dataAspectRatio, axis = 1, where=dataRadius>0)
    stdAspectRatio  = np.std(dataAspectRatio, axis = 1, where=dataRadius>0)

    x = np.flip(np.arange(maxOrder) + 1)

    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col')
    fig.set_figheight(7.2)
    fig.set_figwidth(8.8)
    scale='linear'
    ax1.errorbar(x, 2*meanRadius, stdRadius*2, marker='s', fillstyle='full')
    ax1.set_yscale('log')
    ax1.set_title('Diameter vs. Strahler order')
    ax2.errorbar(x, meanLength, stdLength, marker='s')
    ax2.set_yscale(scale)
    ax2.set_title('Length vs. Strahler order')
    ax3.errorbar(x, meanFlow, stdFlow, marker='s')
    ax3.set_yscale(scale)
    ax3.set_title('Flow vs. Strahler order')
    ax4.errorbar(x, meanAspectRatio, stdAspectRatio, marker='s')
    ax4.set_yscale(scale)
    ax4.set_title('Aspect ratio vs. Strahler order')

    fig2, axes = plt.subplots(2,2, sharex='col')
    fig2.set_figheight(7.2)
    fig2.set_figwidth(8.8)
    scale = 'log'
    axes[0,0].boxplot(dataRadius.transpose())
    axes[0,0].set_yscale(scale)
    axes[0,0].set_title('Radius vs. Strahler order')
    axes[0,1].boxplot(dataLength.transpose())
    axes[0,1].set_yscale(scale)
    axes[0,1].set_title('Length vs. Strahler order')
    axes[1,0].boxplot(dataFlow.transpose())
    axes[1,0].set_yscale(scale)
    axes[1,0].set_title('Flow vs. Strahler order')
    axes[1,1].boxplot(dataAspectRatio.transpose())
    axes[1,1].set_yscale(scale)
    axes[1,1].set_title('Aspect ratio vs. Strahler order')

    if outputImageFile:
        print('\nWriting outputs (image+data) in', outputImageFile[:-4] + '.png/.dat')
        data = np.vstack((x, meanRadius, stdRadius, meanLength, stdLength, meanVolume, stdVolume, meanAspectRatio, stdAspectRatio, orderDistribution))
        np.savetxt(outputImageFile[:-4] + '.dat', data)
        plt.savefig(outputImageFile)
    plt.show()

    return

def BifurcationDiameterRatio(treeData, treeConnectivity, plot=False):
    '''
    Return the ratio of the smaller daughter vessel to the larger (dsdl)
    against the ratio of the larger daughter to the parent vessel (dldp)
    '''

    copyData = copy.deepcopy(treeData)
    copyData.sort()
    data = []
    for Id, parentId, children in treeConnectivity:
        if parentId>0 and len(children)==2: # parentId=-1 if its root

            dp = copyData[Id][1]
            dchildren = [copyData[children[0]][1], copyData[children[1]][1]]
            ds = min(dchildren)
            dl = max(dchildren)

            # dsdl.append(ds/dl)
            # dldp.append(dl/dp)
            data.append([ds/dl, dl/dp, parentId])

    data.sort()
    data = np.array(data)
    dsdl, dldp, Ind = data[:,0], data[:,1], data[:,2]

    if plot:
        aggregatedData = []
        # x = np.linspace(data[:,0].min(),1, 10)
        x = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
        for i,xi in enumerate(x[1:]):
            xim = x[i]
            mask = (dsdl>xim) & (dsdl<xi)
            aggregatedData.append(dldp[mask])

        plt.boxplot(aggregatedData, positions=x[1:])
        # plt.scatter(dsdl, dldp, label='This work')

        m = 2.85
        x = np.linspace(0.2,1,500)
        y = np.linspace(0.2,1,500)
        X,Y = np.meshgrid(x,y)

        f,g = Y**(-m), 1+X**m
        z = f-g
        cs1 = plt.contour(X,Y, z, [0], linestyles='--', linewidths=1.)
        m = 3
        f,g = Y**(-m), 1+X**m
        cs2 = plt.contour(X,Y,(f-g), [0], linestyles='-.', linewidths=1.)

        cs1.collections[0].set_label('m=2.85')
        cs2.collections[0].set_label('m=3')

        plt.xlim(0,1.2)
        # plt.ylim(0,1.2)
        # plt.yscale('log')

        plt.legend()
        plt.show()

    return [a for a in dsdl], [a for a in dldp]

def VBC(treeData, treeConnectivity, plot=False):

    # Copy the radii, sorted by vessel ID
    radii = np.zeros((len(treeData),))
    for v in treeData:
        radii[v[0]] = v[1]

    VBC = []
    for v in treeConnectivity:
        if len(v[-1])==2:
            p, d1, d2 = radii[v[0]], radii[v[-1][0]], radii[v[-1][1]]
            VBC.append( (d1**2 + d2**2)/(p**2) )

    VBC = np.array(VBC)

    if plot:
        split = 10
        x = np.linspace(VBC.min(), VBC.max(), split+1)
        y = np.zeros((split+1,))
        sortedVBC = np.sort(VBC)

        j = 0
        for i in range(split):
            count = 0
            xi = x[i]
            while xi >= sortedVBC[j]:
                count+=1
                j+=1
            y[i] = count
        y[-1] = VBC.size-j

        plt.bar(x[:-1],y[:-1], width=(x[1]-x[0])/2.0)
        plt.xticks(x[:-1], rotation=20)
        plt.xlabel('Vessel branching coefficient')
        plt.ylabel('Number of vessels')
        plt.show()

    return VBC


import sys
from multiprocessing import Process
from multiprocessing.managers import SyncManager
from typing import Any, Callable, List, Optional, Tuple

import boto3
import requests
from flask import Flask, render_template

from rgbserver.config import Config
from rgbserver.control import RGBController

class Webserver(Flask):
    def __init__(self, location: str) -> None:
        super().__init__(__name__)
        self.rgb_config = Config(location)
        self.rgb = RGBController(self.rgb_config)
        self.proc = Process()
        self.manager = SyncManager()
        self.manager.start()
        self.state: List[Tuple[int, int, int]] = self.manager.list()

        self.before_request(self.stop_process)
        self.add_url_rule('/', view_func=self.index)
        self.add_url_rule('/clear', view_func=self.clear)
        self.add_url_rule('/fire', view_func=self.fire)
        self.add_url_rule('/rainbow', view_func=self.rainbow)
        self.add_url_rule('/rainbow_color_wipe', view_func=self.rainbow_color_wipe)
        self.add_url_rule('/rainbow_fade', view_func=self.rainbow_fade)
        self.add_url_rule('/random_fade', view_func=self.random_fade)
        self.add_url_rule('/snake_color', view_func=self.snake_color)
        self.add_url_rule('/snake_fade', view_func=self.snake_fade)
        self.add_url_rule('/snake_rainbow', view_func=self.snake_rainbow)
        self.add_url_rule('/static_color/<int:red>/<int:green>/<int:blue>', view_func=self.static_color)
        self.add_url_rule('/static_color_name/<name>', view_func=self.static_color_name)
        self.add_url_rule('/static_gradient/<int:r1>/<int:g1>/<int:b1>/<int:r2>/<int:g2>/<int:b2>',
                          view_func=self.static_gradient)
        self.add_url_rule('/strobe', view_func=self.strobe)

    def run(self, host: str = None, port: int = None, debug: bool = None, load_dotenv: bool = True,
            **options: Optional[Any]) -> None:
        if self.rgb_config.sqs_url:
            reader = Process(target=self.sqs_reader, args=())
            reader.start()

        super().run('0.0.0.0', self.rgb_config.port, debug, load_dotenv, **options)

    @staticmethod
    def index() -> str:
        return render_template('index.html')

    def clear(self) -> str:
        self.start_process(self.rgb.clear)
        return 'clear'

    def fire(self) -> str:
        self.start_process(self.rgb.fire)
        return 'fire'

    def rainbow(self) -> str:
        self.start_process(self.rgb.rainbow)
        return 'rainbow'

    def rainbow_color_wipe(self) -> str:
        self.start_process(self.rgb.rainbow_color_wipe)
        return 'rainbow_color_wipe'

    def rainbow_fade(self) -> str:
        self.start_process(self.rgb.rainbow_fade)
        return 'rainbow_fade'

    def random_fade(self) -> str:
        self.start_process(self.rgb.random_fade)
        return 'random_fade'

    def snake_color(self) -> str:
        self.start_process(self.rgb.snake, ("color",))
        return 'snake_color'

    def snake_fade(self) -> str:
        self.start_process(self.rgb.snake, ("fade",))
        return 'snake_fade'

    def snake_rainbow(self) -> str:
        self.start_process(self.rgb.snake, ("rainbow",))
        return 'snake_rainbow'

    def static_color(self, red: int, green: int, blue: int) -> str:
        self.start_process(self.rgb.static_color, (red, green, blue))
        return 'static_color'

    def static_color_name(self, name: str) -> str:
        self.start_process(self.rgb.static_color_name, (name,))
        return 'static_color_name'

    def static_gradient(self, r1: int, g1: int, b1: int, r2: int, g2: int, b2: int) -> str:
        self.start_process(self.rgb.static_gradient, ((r1, g1, b1), (r2, g2, b2)))
        return 'static_gradient'

    def strobe(self) -> str:
        self.start_process(self.rgb.strobe)
        return 'strobe'

    def start_process(self, func: Callable, args: Tuple = ()) -> None:
        self.proc = Process(target=func, args=(self.state,) + args)
        self.proc.start()

    def stop_process(self) -> None:
        if self.proc.pid is not None:
            self.proc.terminate()
            self.proc.join()

    def sqs_reader(self) -> None:
        client = boto3.client('sqs', region_name=self.rgb_config.aws_region,
                              aws_access_key_id=self.rgb_config.aws_key,
                              aws_secret_access_key=self.rgb_config.aws_secret)
        while True:
            response = client.receive_message(QueueUrl=self.rgb_config.sqs_url, AttributeNames=[],
                                              MessageAttributeNames=[],
                                              MaxNumberOfMessages=1, VisibilityTimeout=30, WaitTimeSeconds=20)
            if 'Messages' in response:
                message = response['Messages'][0]
                requests.get("http://localhost:" + str(self.rgb_config.port) + "/" + message['Body'])
                client.delete_message(QueueUrl=self.rgb_config.sqs_url, ReceiptHandle=message['ReceiptHandle'])

def main() -> None:
    if len(sys.argv) != 2:
        exit("The only parameter needed is the json config file's absolute path!")
    Webserver(sys.argv[1]).run()

if __name__ == "__main__":
    main()


import pygame as pg
from abc import ABC, abstractmethod

class Drawable(pg.sprite.Sprite, ABC):
    """ The abstract class Drawable represents every object which can be scaled
    and drawn on the map on a specific position. """

    def __init__(self):
        pg.sprite.Sprite.__init__(self)

    @abstractmethod
    def draw(self, surface, scale, pos_offset):
        """ Draws the object on the surface. The object is scaled
        and translated according to the scale and pos_offset parameters. """
        pass

    def _draw_filled_aacircle(self, surface, center, radius, color):
        """ Draws a filled anti-aliased circle. """

        # draw anti-aliasing circle
        pg.gfxdraw.aacircle(
            surface,
            int(center.x), int(center.y),
            radius, color
        )

        # draw filled circle
        pg.gfxdraw.filled_circle(
            surface,
            int(center.x), int(center.y),
            radius, color
        )

    def _draw_filled_aatriangle(self, surface, p1, p2, p3, color):
        """ Draws a filled anti-aliased triangle. """

        pg.gfxdraw.filled_trigon(
            surface,
            int(p1.x), int(p1.y),
            int(p2.x), int(p2.y),
            int(p3.x), int(p3.y),
            color,
        )

        pg.gfxdraw.aatrigon(
            surface,
            int(p1.x), int(p1.y),
            int(p2.x), int(p2.y),
            int(p3.x), int(p3.y),
            color,
        )


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

# Learn more: https://github.com/sequgen/accelsim

from setuptools import setup, find_packages

with open('README.md') as f:
    readme = f.read()

#with open('LICENSE') as f:
#    license = f.read()

setup(
    name='accelsim',
    version='0.1.0',
    description="Simulate accelerometer's readings",
    long_description=readme,
    author='Pablo Rodríguez-Sánchez',
    author_email='pablo.rodriguez.sanchez@gmail.com',
    url='https://github.com/sequgen/accelsim',
    # license=license,
    install_requires=[
        'numpy',
    ],
    packages=find_packages(
        exclude=('tests', 'docs', 'vignettes', 'scripts', 'drafts')
        )
)


import os

import psycopg2
import pytest

from ddtrace.contrib.psycopg.patch import patch
from ddtrace.contrib.psycopg.patch import unpatch
from ddtrace.vendor import wrapt
from tests.contrib.config import POSTGRES_CONFIG
from tests.utils import override_config

@pytest.fixture(autouse=True)
def patch_psycopg():
    patch()
    assert isinstance(psycopg2.connect, wrapt.ObjectProxy)
    yield
    unpatch()

@pytest.mark.snapshot()
def test_connect_default():
    """By default we do not trace psycopg2.connect method"""
    conn = psycopg2.connect(**POSTGRES_CONFIG)
    assert conn

@pytest.mark.snapshot(wait_for_num_traces=1)
def test_connect_traced():
    """When explicitly enabled, we trace psycopg2.connect method"""
    with override_config("psycopg", {"trace_connect": True}):
        conn = psycopg2.connect(**POSTGRES_CONFIG)
        assert conn

@pytest.mark.snapshot(token="tests.contrib.psycopg2.test_psycopg_snapshot.test_connect_traced", wait_for_num_traces=1)
def test_connect_traced_via_env(run_python_code_in_subprocess):
    """When explicitly enabled, we trace psycopg2.connect method"""

    code = """
import psycopg2

import ddtrace
from tests.contrib.config import POSTGRES_CONFIG

ddtrace.patch_all()

conn = psycopg2.connect(**POSTGRES_CONFIG)
assert conn
    """

    env = os.environ.copy()
    env["DD_PSYCOPG2_TRACE_CONNECT"] = "true"
    out, err, status, pid = run_python_code_in_subprocess(code, env=env)
    assert status == 0, err
    assert out == b"", err
    assert err == b""


'''
Created on 20 mar. 2019

@author: julio
'''
import socketserver
import threading
import django
from queue import Queue, Empty

django.setup()

from fgserver.models import Airport, Order
from fgserver.messages import PosMsg
from fgserver import llogger, settings, setInterval
from fgserver.server.fgmpie import pie_msg, PacketData
from fgserver.signals import message_received, signal_server_started
from fgserver.server.utils import get_pos_msg, process_msg

class FGServer(socketserver.ThreadingMixIn, socketserver.UDPServer):
    incoming = Queue()
    outgoing = Queue()
    _thread = None

    def init(self):
        self._thread = threading.Thread(target=self.serve_forever)
        self._thread.daemon = True
        server_thread.start()

    def put(self,msg):
        self.outgoing.put(msg)

    def get(self,msg):
        self.incoming.get(msg)

class FGHandler(socketserver.BaseRequestHandler):

    def handle(self):
        data = self.request[0]
        try:
            unp = PacketData(data)
            pos = PosMsg()
            pos.receive_pie(unp)
            pos.header.reply_addr=self.client_address[0]
            pos.header.reply_port=self.client_address[1]
            #print("R",pos)
            #process_msg(pos)
            self.server.incoming.put(pos)
        except:
            llogger.exception("While receiving message")

def startup_tasks():
    Order.objects.all().update(expired=True)

@setInterval(0.5)
def send_msg():
    for airport in Airport.objects.filter(active=True):
        msg = get_pos_msg(airport)
        #llogger.debug("Sending pos for AI ATC   %s" % msg)
        server.outgoing.put(msg)

#     from fgserver.ai.models import FlightPlan
#     for plan in FlightPlan.objects.filter(enabled=True):
#         try:
#             pos = plan.aircraft.status.get_position_message()
#             server.outgoing.put(pos)
#         except AircraftStatus.DoesNotExist:
#             AircraftStatus(aircraft=plan.aircraft, date=timezone.now()).save()
#             pos = plan.aircraft.status.get_position_message()
#             server.outgoing.put(pos)
#         except:
#             llogger.exception("In loop")

#     else:
#         diff =(timezone.now() - (aircraft.updated or timezone.now())).total_seconds()
#             if diff < 2:
#                 #print("saving",aircraft.__dict__)
#                 aircraft.save()

if __name__ == "__main__":
    HOST, PORT = "0.0.0.0", 5100

    server = FGServer((HOST, PORT), FGHandler)
    server_thread = threading.Thread(target=server.serve_forever)
    server_thread.daemon = True

    startup_tasks()

    try:
        server_thread.start()
        #signal_server_started.send_robust(None)
        print("Server started at {} port {}".format(HOST, PORT))
        send_msg()
        while True:
            try:
                pos = server.incoming.get(True, .1)
                #llogger.debug("Received %s" % pos )
                message_received.send_robust(sender=server,msg=pos)
                process_msg(pos)

            except Empty:
                pass
            try:
                msg = server.outgoing.get(False)
                #llogger.debug("Sending message %s" % msg )
                buff = pie_msg(msg)
                server.socket.sendto(buff,settings.FGATC_RELAY_SERVER)
            except Empty:
                pass

    except (KeyboardInterrupt, SystemExit):
        server.shutdown()
        server.server_close()
        exit()


import torch
import torch.nn as nn
import torch.nn.functional as F

class UpSampleBN(nn.Module):
    def __init__(self, skip_input, output_features):
        super(UpSampleBN, self).__init__()

        self._net = nn.Sequential(nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.BatchNorm2d(output_features),
                                  nn.LeakyReLU(),
                                  nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1),
                                  nn.BatchNorm2d(output_features),
                                  nn.LeakyReLU())

    def forward(self, x, concat_with):
        up_x = F.interpolate(x, size=[concat_with.size(2), concat_with.size(3)], mode='bilinear', align_corners=True)
        f = torch.cat([up_x, concat_with], dim=1)
        return self._net(f)

class DecoderBN(nn.Module):
    def __init__(self, num_features=2048, num_classes=1, bottleneck_features=2048):
        super(DecoderBN, self).__init__()
        features = int(num_features)

        self.conv2 = nn.Conv2d(bottleneck_features, features, kernel_size=1, stride=1, padding=1)

        self.up1 = UpSampleBN(skip_input=features // 1 + 112 + 64, output_features=features // 2)
        self.up2 = UpSampleBN(skip_input=features // 2 + 40 + 24, output_features=features // 4)
        self.up3 = UpSampleBN(skip_input=features // 4 + 24 + 16, output_features=features // 8)
        self.up4 = UpSampleBN(skip_input=features // 8 + 16 + 8, output_features=features // 16)

        #         self.up5 = UpSample(skip_input=features // 16 + 3, output_features=features//16)
        self.conv3 = nn.Conv2d(features // 16, num_classes, kernel_size=3, stride=1, padding=1)
        # self.act_out = nn.Softmax(dim=1) if output_activation == 'softmax' else nn.Identity()

    def forward(self, features):
        x_block0, x_block1, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[
            11]

        x_d0 = self.conv2(x_block4)

        x_d1 = self.up1(x_d0, x_block3)
        x_d2 = self.up2(x_d1, x_block2)
        x_d3 = self.up3(x_d2, x_block1)
        x_d4 = self.up4(x_d3, x_block0)
        x_d5 = self.up5(x_d4, features[0])
        out = self.conv3(x_d5)
        # out = self.act_out(out)
        # if with_features:
        #     return out, features[-1]
        # elif with_intermediate:
        #     return out, [x_block0, x_block1, x_block2, x_block3, x_block4, x_d1, x_d2, x_d3, x_d4]
        return out

class Encoder(nn.Module):
    def __init__(self, backend):
        super(Encoder, self).__init__()
        self.original_model = backend

    def forward(self, x):
        features = [x]
        for k, v in self.original_model._modules.items():
            if (k == 'blocks'):
                for ki, vi in v._modules.items():
                    features.append(vi(features[-1]))
            else:
                features.append(v(features[-1]))
        return features

class UEfficientNetEncoder(nn.Module):
    def __init__(self, backend):
        super(UEfficientNetEncoder, self).__init__()
        self.encoder = Encoder(backend)
        self.decoder = DecoderBN(num_classes=128)

    def forward(self, x, **kwargs):
        unet_out = self.decoder(self.encoder(x), **kwargs)
        return unet_out

    def get_1x_lr_params(self):  # lr/10 learning rate
        return self.encoder.parameters()

    def get_10x_lr_params(self):  # lr learning rate
        modules = [self.decoder]
        for m in modules:
            yield from m.parameters()

    @classmethod
    def build(cls, n_bins, **kwargs):
        basemodel_name = 'tf_efficientnet_b5_ap'

        print('Loading base model ()...'.format(basemodel_name), end='')
        basemodel = torch.hub.load('rwightman/gen-efficientnet-pytorch', basemodel_name, pretrained=True)
        print('Done.')

        # Remove last layer
        print('Removing last two layers (global_pool & classifier).')
        basemodel.global_pool = nn.Identity()
        basemodel.classifier = nn.Identity()

        # Building Encoder-Decoder model
        print('Building Encoder-Decoder model..', end='')
        m = cls(basemodel, **kwargs)
        print('Done.')
        return m

if __name__ == '__main__':
    model = UEfficientNetEncoder.build(100)
    x = torch.rand(2, 3, 480, 640)
    pred = model(x)
    print(pred.shape)


import pandas as pd
import matplotlib.pyplot as plt
import sys
import numpy as np
import sklearn as sk
import seaborn as sns
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV
from joblib import dump
from joblib import load
from sklearn.feature_selection import RFE
from sklearn.metrics import f1_score
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

def build_pipeline(classifier, isScaled, selectFeatures):
    """
    This function takes a scikit-learn classifier and builds a ML
    pipeline that applies a series of transformations.

    Parameters:
    - classifier: The machine learning classifer of interest ex. KNN
    - isScaled(bool): Should we scale data? True or False
    - selectFeatures(bool): Should we perform RFE? True or False

    Returns:
    - pipe(Pipeline): A ML pipeline capable of transforming data

    """
    try:
        if selectFeatures:
            if isScaled:
                pipe = Pipeline(steps=[("imputer", SimpleImputer()),
                                       ("scale", StandardScaler()),
                                       ("rfe", RFE(estimator=classifier)),
                                       ("clf", classifier)])
            else:
                pipe = Pipeline(steps=[("imputer", SimpleImputer()),
                                       ("rfe", RFE(estimator=classifier)),
                                       ("clf", classifier)])
        else:
            if isScaled:
                pipe = Pipeline(steps=[("imputer", SimpleImputer()),
                                       ("scale", StandardScaler()),
                                       ("clf", classifier)])
            else:
                pipe = Pipeline(steps=[("imputer", SimpleImputer()),
                                       ("clf", classifier)])
        return pipe

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def run_gridsearch(pipe, param_grid, num_folds, metric,
                   X_train, y_train, output_model):
    """
    This function runs a grid search over a combination of
    hyperparameters using k-fold cross validation to select the best
    estimator

    Parameters:
    - pipe(Pipeline): A ML pipeline capable of transforming data
    - param_grid(dict): The desired hyperparameters and ranges to tune
    - num_folds(int): The number of folds for cross validation
    - metric(str): The evaluation metric  ex. "f1"
    - X_train(ndarray): A numpy array containing 1-test_size fraction of
                        the original dataset
    - y_train(ndarray): A numpy array containing class labels for X_train
    - output_model: The name of the output file containing the best model

    Returns:
    - best_params_(dict): The best performing hyperparameters
    - best_score_(float): The average validation score over k-folds
    - best_estimator_: The actual best performing model

    """

    try:
        search = GridSearchCV(pipe, param_grid, cv=num_folds,
                              scoring=metric, n_jobs=-1,
                              return_train_score=True)
        search.fit(X_train, y_train)
        dump(search.best_estimator_, output_model+".pkl")
        return search.best_params_, search.best_score_, search.best_estimator_, search.cv_results_  # noqa: E501

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def import_data(csv_file):
    """
    This function imports data from a .csv file as a pandas dataframe

    Parameters:
    - csv_file(str): A .csv file path

    Returns:
    - df(DataFrame): A pandas DataFrame containing features and class labels

    """
    try:
        df = pd.read_csv(csv_file)
        return df

    except FileNotFoundError as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def make_correlation_heatmap(df, title, outfile):
    """
    This function creates a correlation heatmap for a dataframe
    of features and saves it to a user-defined file.

    Parameters:
    - df(DataFrame): A dataframe of X features and y label
    - title: The title of the correlation heat map
    - outfile(str): The file where the heat map will be saved

    Returns:
    - None, but a .png file is created

    """

    try:
        df_corr = df.corr()
        plt.figure(figsize=(10, 10))
        plt.title(title)
        mask = np.zeros_like(df_corr)
        mask[np.triu_indices_from(mask)] = True
        with sns.axes_style("white"):
            ax = sns.heatmap(df_corr, fmt=".2f", mask=mask, annot=True)
        bottom, top = ax.get_ylim()
        ax.set_ylim(bottom + 0.5, top - 0.5)
        plt.savefig(outfile)

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def one_hot_encode(df, col_label):
    """
    This function one-hot encodes a single categorical feature in a DataFrame
    and drops the original feature column from the DataFrame

    Parameters:
    - df(DataFrame): A pandas DataFrame containing features and class labels
    - col_label(str): The column label of categorical feature

    Returns:
    - df(DataFrame): A pandas DataFrame after one-hot encoding

    """
    try:
        df_enc = pd.get_dummies(df[col_label])
        df = pd.concat([df_enc, df], axis=1)
        df.drop([col_label], axis=1, inplace=True)

    except KeyError as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

    return df

def split_data(df, test_size, label_col):
    """
    This function splits a pandas DataFrame into X_train, X_test, y_train, and
    y_test sets according to user-defined percentage.

    Parameters:
    - df(DataFrame): A pandas DataFrame containing features and class labels
    - test_size(float): The fraction of original DataFrame held out for testing
                        Ex. 0.2
    - label_col(str): The column of the DataFrame that contains the class label

    Returns:
    - X_train(ndarray): A numpy array containing 1-test_size fraction of
                        the original dataset
    - y_train(ndarray): A numpy array containing the class labels for X_train
    - X_test(ndarray): A numpy array containing test_size fraction of
                       the original dataset
    - y_test(ndarray): A numpy array containing the class labels for X_test

    """
    try:
        train, test = train_test_split(df, test_size=test_size)
        X_train = train.loc[:, train.columns != label_col].values
        y_train = train[label_col].values
        X_test = test.loc[:, train.columns != label_col].values
        y_test = test[label_col].values

        return X_train, y_train, X_test, y_test

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def make_confusion_matrix(y_true, y_pred, output_file):
    """
    This function creates a confusion matrix in the form of a heat map

    Parameters:
    - y_true(list): The true class labels for X_test
    - y_pred(list): The predicted class labels for X_test
    - output_file(str): The file path for saving the confusion matrix

    Returns:
    - None; a .png file of the confusion matrix is created

    """
    try:
        data = confusion_matrix(y_true, y_pred)
        df_cm = pd.DataFrame(data, columns=np.unique(y_true),
                             index=np.unique(y_true))
        df_cm.index.name = "Actual"
        df_cm.columns.name = "Predicted"
        plt.figure(figsize=(10, 10))
        plt.title("Confusion matrix")
        cm = sns.heatmap(df_cm, annot=True)
        plt.savefig(output_file)

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)

def evaluate_test_set(y_test, y_pred, output_file):
    """
    This function creates a confusion matrix of the actual
    vs. predicted class labels and returns the f1_score

    Parameters:
    - y_test(ndarray): A numpy array containing the class labels for X_test
    - y_test(ndarray): A numpy array containing the predicted
                       class labels for X_test
    - output_file(str): The file path for saving the confusion matrix

    Returns:
    - score(float): The f1_score for the model's predictions on the test set

    """
    try:
        make_confusion_matrix(y_test, y_pred, output_file)
        return f1_score(y_test, y_pred)

    except Exception as inst:
        print("Run-Time Error:", type(inst))
        sys.exit(1)