# -*- coding: utf-8 -*-
import json
import random
import time

import scrapy

from ljh.items import *

class DoubanSpider(scrapy.Spider):
    name = 'douban'
    allowed_domains = ['movie.douban.com']
    mapping = {'姓名': 'name', '性别': 'gender', '星座': 'constellation', '出生日期': 'birthday', '去世日期': 'date_of_death',
               '出生地': 'birthplace',
               '职业': 'job', '更多外文名': 'other_foreign_names', '更多中文名': 'other_chinese_names', '生卒日期': 'birthday',
               '家庭成员': 'families'}

    # start_urls = ['https://movie.douban.com/top250']
    # start_urls = ['https://movie.douban.com/subject/30304469/?tag=%E7%83%AD%E9%97%A8&from=gaia']
    count = 0
    maxcount = 2500

    def start_requests(self):
        # yield scrapy.Request(url='https://movie.douban.com/subject/1292052/', callback=self.parse_film_page)
        # yield scrapy.Request(url='https://movie.douban.com/top250', callback=self.parse)
        tags = [
            '%e7%83%ad%e9%97%a8',  # 热门
            '%e6%9c%80%e6%96%b0',  # 最新
            '%e7%bb%8f%e5%85%b8',  # 经典
            '%e8%b1%86%e7%93%a3%e9%ab%98%e5%88%86'  # 豆瓣高分
        ]
        for tag in tags:
            yield scrapy.Request(
                url='https://movie.douban.com/j/search_subjects?type=movie&tag=%s&sort=recommend&page_limit=2000&page_start=0' % tag,
                callback=self.parse_json)

    def parse_json(self, response):
        res = json.loads(response.body_as_unicode())['subjects']
        for i in res:
            yield response.follow(i['url'], callback=self.parse_film_page)

    # def start_requests(self):
    #     # yield scrapy.Request(url='https://movie.douban.com/subject/1292052/', callback=self.parse_film_page)
    #     yield scrapy.Request(url='https://movie.douban.com/top250', callback=self.parse)
    #
    # def parse(self, response):
    #     res = response.css('ol.grid_view li div.item div.pic a::attr(href)').getall()
    #     for url in res:
    #         yield response.follow(url, callback=self.parse_film_page)
    #         # print(self.count, picture, main_page, title, description, star, quote)
    #     next = response.css('div.paginator span.next a::attr(href)').extract_first()
    #     if next is not None:
    #         yield response.follow(next, callback=self.parse)

    def parse_film_page(self, response):
        self.count = self.count + 1
        item = FilmItem()
        item['film_id'] = response.request.url.split('subject/')[1].split('/')[0]
        item['title'] = response.css('#content > h1 > span:nth-child(1)::text').extract_first()
        item['year'] = response.css('#content > h1 > span.year::text').extract_first().replace('(', '').replace(')', '')
        item['directors'] = response.css('#info > span:nth-child(1) > span.attrs > a::text').getall()
        director_hrefs = response.css('#info > span:nth-child(1) > span.attrs > a::attr(href)').getall()
        item['writers'] = response.css('#info > span:nth-child(3) > span.attrs > a::text').getall()
        writer_hrefs = response.css('#info > span:nth-child(3) > span.attrs > a::attr(href)').getall()
        item['actors'] = response.css('#info > span.actor > span.attrs > a::text').getall()[:3]
        actor_hrefs = response.css('#info > span.actor > span.attrs > a::attr(href)').getall()[:3]
        item['types'] = response.css('#info > span[property="v:genre"]::text').getall()
        item['region'] = \
            response.css('#info').re(r'制片国家/地区:.*
')[0].split('')[1].split('
')[0]
        item['description'] = response.css('#link-report > span::text').extract_first().strip()
        try:
            item['language'] = response \
                .css('#info') \
                .re(r'语言:.*
')[0] \
                .split('')[1] \
                .split('
')[0] \
                .strip()
        except IndexError:
            item['language'] = ''

        try:
            item['other_name'] = response \
                .css('#info') \
                .re(r'又名:.*
')[0] \
                .split('')[1] \
                .split('
')[0] \
                .strip()
        except IndexError:
            item['other_name'] = ''
        item['release_date'] = response.css('#info > span[property="v:initialReleaseDate"]::text').extract_first()
        item['size'] = response.css('#info > span[property="v:runtime"]::attr(content)').extract_first()
        item['star'] = response.css(
            '#interest_sectl > div.rating_wrap.clearbox > div.rating_self.clearfix > strong::text').extract_first()
        item['star_percent'] = response.css(
            '#interest_sectl > div.rating_wrap.clearbox > div.ratings-on-weight > div > span.rating_per::text').getall()
        item['comment_num'] = response.css(
            '#interest_sectl > div.rating_wrap.clearbox > div.rating_self.clearfix > div > div.rating_sum > a > span::text').extract_first()
        item['src'] = response.css(
            '#mainpic > a > img::attr(src)').extract_first()
        recommendations = response.css('#recommendations > div > dl > dt > a::attr(href)').getall()
        # print(title, year, directors, writers, actors, types, region, release_date, size, star, star_percent, sep='\n')
        try:
            for href in director_hrefs:
                d_item = DirectorItem()
                d_item['film_id'] = item['film_id']
                d_item['celebrity_id'] = href.split('celebrity/')[1].split('/')[0]
                yield d_item
            for href in writer_hrefs:
                w_item = WriterItem()
                w_item['film_id'] = item['film_id']
                w_item['celebrity_id'] = href.split('celebrity/')[1].split('/')[0]
                yield w_item
            for href in actor_hrefs:
                a_item = ActorItem()
                a_item['film_id'] = item['film_id']
                a_item['celebrity_id'] = href.split('celebrity/')[1].split('/')[0]
                yield a_item
        except IndexError:
            pass
        else:
            yield item
        for i in list(director_hrefs) + writer_hrefs + actor_hrefs:
            yield response.follow(i, callback=self.parse_celebrity_page)

        if self.count < self.maxcount:
            for i in recommendations:
                yield response.follow(i, callback=self.parse_film_page)

    def parse_celebrity_page(self, response):
        name = response.css('#content > h1::text').extract_first()
        ul = response.css('#headline > div.info > ul > li::text').getall()
        info_names = ['姓名'] + response.css('#headline > div.info > ul > li > span::text').getall()
        del info_names[-1]
        info_list = [name]
        for s in ul:
            s = s.replace(':', '').strip()
            if len(s) > 0:
                info_list.append(s)
        # result = {k: v for k, v in zip(info_names, info_list)}
        # print(result)
        item = CelebrityItem()
        item['celebrity_id'] = response.request.url.split('celebrity/')[1].split('/')[0]
        item['human_src'] = response.css('#headline > div.pic > a > img::attr(src)').extract_first()
        for k in self.mapping:
            item[self.mapping[k]] = ''
        for k, v in zip(info_names, info_list):
            if k == '生卒日期':
                birthday_and_date_of_death = v.split('至')
                item['birthday'] = birthday_and_date_of_death[0]
                item['date_of_death'] = birthday_and_date_of_death[1]
            else:
                item[self.mapping[k]] = v
        yield item


import json
import os
import os.path as osp
import pickle
from collections import Counter

import numpy as np
import torch
import torch.utils.data as data
from PIL import Image

from utils import text
from utils.image import coco_name_format

def collate_fn(batch):
    """Sort batch (list) on question lengths for use in RNN pack_padded_sequence later."""
    batch.sort(key=lambda x: x['question_len'], reverse=True)
    return data.dataloader.default_collate(batch)

def get_dataloader(annotations, questions, images, args,
                   split="train", maps=None, vocab=None,
                   raw_images=False, transforms=None, shuffle=True):
    return data.DataLoader(VQADataset(annotations, questions, images, split, args,
                                      raw_images=raw_images, vocab=vocab,
                                      transforms=transforms, maps=maps),
                           batch_size=args.batch_size,
                           num_workers=args.num_workers,
                           shuffle=shuffle,
                           collate_fn=collate_fn)

class VQADataset(data.Dataset):
    def __init__(self, annotations, questions, images_dataset, split, args,
                 raw_images=False, transforms=None,
                 vocab=None, normalize_img=True, maps=None, year=2014):

        # the data is saved as a dict where the key is the image_id
        # and the value is the VGG feature vector
        if not raw_images:
            self.images_dataset = torch.load(images_dataset)
            print("Loaded {0} image embeddings dataset".format(split))

        self.split = split
        self.year = year

        self._process_dataset(annotations, questions, args, split, maps=maps)

        if vocab:
            self.vocab = vocab

        self.embed_question = args.embed_question

        self.raw_images = raw_images
        self.normalize_img = normalize_img
        self.transforms = transforms
        self.root = args.image_root

    def _process_dataset(self, annotations, questions, args,  split="train", maps=None):
        """
        Process the dataset and load it up.
        We should only do this for the training set.
        """
        self.data, self.vocab, self.word_to_wid, self.wid_to_word, \
            self.ans_to_aid, self.aid_to_ans = \
            process_vqa_dataset(questions, annotations, split,
                                maps, args.top_answer_limit, args.max_length)

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

    def __getitem__(self, index):
        d = self.data[index]

        item = dict()

        if self.raw_images:
            item["image_id"] = d["image_id"]
            img = Image.open(osp.join(osp.expanduser(self.root),
                                      "{0}{1}".format(self.split, self.year),
                                      d["image_name"]))
            img = img.convert(mode='RGB')

            if self.transforms is not None:
                img = self.transforms(img)

        else:
            # Process Visual (image or features)
            # the preprocess script should have already saved these as Torch tensors
            item["image_id"] = d["image_id"]
            img = self.images_dataset[d["image_id"]].squeeze()

        item['image'] = img

        # Process Question (word token)
        item['question_id'] = d['question_id']
        item['question_wids'] = d['question_wids'].astype(np.int64)

        # if self.embed_question:
        #     item['question'] = torch.from_numpy(d['question_wids'])
        # else:
        #     # one_hot_vec = np.zeros((len(d["question_wids"]), len(self.vocab)))
        #     # for k in range(len(d["question_wids"])):
        #     #     one_hot_vec[k, d['question_wids'][k]] = 1
        #     # item['question'] = torch.from_numpy(one_hot_vec).float()
        #
        #     one_hot_vec = torch.zeros(
        #         (len(d["question_wids"]), len(self.vocab)))
        #     one_hot_vec[torch.arange(one_hot_vec.size(0)).long(
        #     ), d['question_wids'].astype(np.int64)] = 1
        #     item['question'] = one_hot_vec.float()
        item['question'] = torch.from_numpy(d['question_wids']).long()

        item['question_len'] = d['question_length']
        item['answer_type'] = d['answer_type']

        if self.split == "train":
            item['answer_id'] = d['answer_id']

        return item

def process_vqa_dataset(questions_file, annotations_file, split, maps=None,
                        top_answer_limit=1000, max_length=26, year=2014):
    """
    Process the questions and annotations into a consolidated dataset.
    This is done only for the training set.
    :param questions_file:
    :param annotations_file:
    :param split: The dataset split.
    :param maps: Dict containing various mappings such as \
        word_to_wid, wid_to_word, ans_to_aid and aid_to_ans.
    :param top_answer_limit:
    :param max_length: The maximum quetsion length. Taken from the VQA sample code.
    :param year: COCO Dataset release year.
    :return: The processed dataset ready to be used

    """
    cache_file = "vqa_{0}_dataset_cache.pickle".format(split)

    # Check if preprocessed cache exists. If yes, load it up, else preprocess the data
    if os.path.exists(cache_file):
        print("Found {0} set cache! Loading...".format(split))
        dataset, vocab, word_to_wid, wid_to_word, ans_to_aid, aid_to_ans = pickle.load(
            open(cache_file, 'rb'))

    else:
        # load the annotations and questions files
        print("Loading {0} annotations".format(split))
        with open(annotations_file) as ann:
            j = json.load(ann)
            annotations = j["annotations"]

        print("Loading {0} questions".format(split))
        with open(questions_file) as q:
            j = json.load(q)
            questions = j["questions"]

        # load up the dataset
        dataset = []
        for idx, q in enumerate(questions):
            d = dict()
            d["question_id"] = q["question_id"]
            d["question"] = q["question"]
            d["image_id"] = q["image_id"]
            d["image_name"] = coco_name_format(q["image_id"], split, year)

            d["answer"] = annotations[idx]["multiple_choice_answer"]
            answers = []
            for ans in annotations[idx]['answers']:
                answers.append(ans['answer'])
            d['answers_occurence'] = Counter(answers).most_common()

            d["question_type"] = annotations[idx]["question_type"]
            d["answer_type"] = annotations[idx]["answer_type"]

            dataset.append(d)

        if split == "train":
            # Get the top N answers so we can filter the dataset
            # to only questions with these answers
            top_answers = text.get_top_answers(dataset, top_answer_limit)
            dataset = text.filter_dataset(dataset, top_answers)

            # Process the questions
            dataset = text.preprocess_questions(dataset)

            vocab = text.get_vocabulary(dataset)

            # 0 is used for padding
            word_to_wid = {w: i+1 for i, w in enumerate(vocab)}
            wid_to_word = {i+1: w for i, w in enumerate(vocab)}
            ans_to_aid = {a: i for i, a in enumerate(top_answers)}
            aid_to_ans = {i: a for i, a in enumerate(top_answers)}

            dataset = text.encode_answers(dataset, ans_to_aid)

        else:  # split == "val":
            # Process the questions
            dataset = text.preprocess_questions(dataset)

            vocab = maps["vocab"]

            word_to_wid = maps["word_to_wid"]
            wid_to_word = maps["wid_to_word"]
            ans_to_aid = maps["ans_to_aid"]
            aid_to_ans = maps["aid_to_ans"]

            dataset = text.remove_tail_words(dataset, vocab)

        dataset = text.encode_questions(dataset, word_to_wid, max_length)

        print("Caching the processed data")
        pickle.dump([dataset, vocab, word_to_wid, wid_to_word, ans_to_aid, aid_to_ans],
                    open(cache_file, 'wb+'))

    return dataset, vocab, word_to_wid, wid_to_word, ans_to_aid, aid_to_ans


"""Low level module for IgBLAST api calls

You probably only want to interact with this module through airr as the database files are extremely tricky to get right

"""
from __future__ import annotations

import glob
import logging
import os
import subprocess
import tempfile
import warnings
from multiprocessing import cpu_count
from pathlib import Path
from typing import Any, List, Union

# Third party
import pandas as pd
import semantic_version

from sadie.airr.airrtable.constants import IGBLAST_AIRR
from sadie.airr.exceptions import (
    BadIgBLASTArgument,
    BadIgBLASTExe,
    BadIgDATA,
    IgBLASTRunTimeError,
    MissingIgBLASTArgument,
)

# package/module level
from sadie.utility.util import is_tool

# get logger in global scope
logger = logging.getLogger("IgBLAST")

def ensure_prefix_to(path: Union[str, Path]) -> Union[Path, bool]:
    """Ensure that the blast db is actually a blast like database

    The problem is that a blast_db takes in a prefix file

    ex. /path/to/blast/human_V

    which is a file or path that does not actually exists, but blasts uses it as a file glob to match
    /path/to/blast/human_V.nod
    /path/to/blast/human_V.nsq
    /path/to/blast/human_V.fasta

    because of this, we don't have validated file path but an actual file path glob. This method validates that the file glob returns blast like files

    Parameters
    ----------
    path : str
        path glob

    Returns
    -------
    return : Union[str, bool]
       returns Path or False if not a file glob
    """

    # convert Path to str
    path = str(path)
    directory_path = os.path.dirname(path)
    if not os.path.exists(directory_path):
        return False

    # get abs path to directory
    directory_path = os.path.abspath(directory_path)

    # base name
    basename = os.path.basename(path)
    glob_path = os.path.join(directory_path, basename) + "*"
    # make sure that there are things that match this queyr
    if not glob.glob(glob_path):
        return False
    return Path(os.path.join(directory_path, basename))

class IgBLASTArgument:
    """
    A class for handling all IgBLAST Arguments
    """

    def __init__(self, name: str, arg_key: str, arg_value: Union[str, int, bool, Path], required: bool):
        """IgBLASTArgument Class constructor

        Parameters
        ----------
        name : str
            the internal name for the computer
        arg_key : str
           the argument key, ex. -germline_db_V
        arg_value : Union[str, int, bool]
           the value for the argument /path/germline/db/V
        required : bool
           is the argument required
        """
        self.name = name
        self.key = arg_key
        self.value = arg_value
        self.required = required

    @property
    def name(self) -> str:
        """
        An internal name for the argument
        """
        return self._name

    @name.setter
    def name(self, n: str) -> None:
        self._name = n

    @property
    def key(self) -> str:
        """The blast command key argument

        ex. '-germline_db_v

        Returns
        -------
        str
           blast command key
        """
        return self._key

    @key.setter
    def key(self, k: str) -> None:
        self._key = k

    @property
    def value(self) -> Union[str, int, bool, Path]:
        """Return the value of the argument

        Returns
        -------
        Union[str,int,bool]
           ex. /path/to/database
        """
        return self._value

    @value.setter
    def value(self, v: Union[str, int, bool]) -> None:
        self._value = v

    @property
    def required(self) -> bool:
        """Returns if the argument is required

        Returns
        -------
        bool
            if argument is required
        """
        return self._required

    @required.setter
    def required(self, r: bool) -> None:
        self._required = r

    def get_formatted_blast_arg(self) -> List[str]:
        """Return the blast formatted argument as a list

        Returns
        -------
        Union[List[str],List[str,str]]
           Either returns a single argument ['-arg'] for bool args or key value arguments ['-arg', 'value']
        """
        # If value is a bool, we return the key
        if isinstance(self.value, bool) and self.value:
            return ["-" + self.key]
        else:
            # If its not a bool, we check if it has been set
            if self.value:
                return ["-" + str(self.key), str(self.value)]
        return []  # return an empty str if its been set

    def __str__(self) -> str:
        return "{}-{}".format(self.name, self.key)

ArgumentType = Union[IgBLASTArgument, int, str, Path]

class IgBLASTN:
    """IgBLASTN

    IgBLASTN class.  A tool for immunoglobulin (IG) and T cell receptor (TR) V domain sequences from nucletodies.

    This is a lower level class and you should probably use sadie.airr to interact

    Examples
    --------
    >>> ig_blast = igblast.IgBLASTN('igblastn')
    >>> germline_ref = "reference/germlines/"
    >>> db_ref = "reference/germlines/blastdb/Ig"
    >>> aux_path = "reference/germlines/aux_data)

    # Set data
    >>> ig_blast.igdata = germline_ref

    >>> query = "fasta_inputs/PG9_H.fasta"
    >>> ig_blast.germline_db_v = os.path.join(db_ref, "human/human_V")
    >>> ig_blast.germline_db_d = os.path.join(db_ref, "human/human_D")
    >>> ig_blast.germline_db_j = os.path.join(db_ref, "human/human_J")
    >>> ig_blast.aux_path = os.path.join(aux_path, "human_gl.aux")
    >>> ig_blast.organism = "human"
    >>> csv_dataframe = ig_blast.run(query)
    """

    # Only allow these attributes
    __slots__ = [
        "_executable",
        "_version",
        "_min_d_match",
        "_num_v",
        "_num_d",
        "_num_j",
        "_outfmt",
        "_receptor",
        "_word_size",
        "_nomenclature",
        "_gap_open",
        "_gap_extend",
        "_num_threads",
        "_extend_5",
        "_extend_3",
        "_j_penalty",
        "_v_penalty",
        "_d_penalty",
        "_organism",
        "_germline_db_v",
        "_germline_db_d",
        "_germline_db_j",
        "_germline_db_c",
        "_aux_path",
        "_igdata",
        "_temp_dir",
        "_allow_vdj_overlap",
    ]

    def __init__(self, executable: Union[Path, str] = "igblastn", tmp_dir: Union[Path, str] = None):
        """IgBLASTN with a query. Set everything up with a setter

        Parameters
        ----------
        excutable : Union[Path,str], optional
            Path to the igblastn executable, by default "igblastn"
        tmp_dir : Union[Path,str], optional
            Path to the temporary directory, by default None
        """

        # set the executable dynamically
        self.executable = Path(executable)
        self._version = self._get_version()

        # setup all the default values if we don't add them
        self.min_d_match = 5
        self.num_v = 3
        self.num_d = 3
        self.num_j = 3
        self.outfmt = 19
        self.receptor = "Ig"
        self.word_size = 5
        self.nomenclature = "imgt"
        self.gap_open = 5
        self.gap_extend = 2
        self.num_threads = cpu_count()
        self.extend_5 = True
        self.extend_3 = True
        self.j_penalty = -2
        self.v_penalty = -1
        self.d_penalty = -1
        self.allow_vdj_overlap = False

        # Make these blank, if they are not set by the caller, then we will complain during runtime. They must be set dynamically
        self._organism = IgBLASTArgument("organism", "organism", "", True)
        self._germline_db_v = IgBLASTArgument("germline_db_v", "germline_db_V", "", True)
        self._germline_db_d = IgBLASTArgument("germline_db_d", "germline_db_D", "", True)
        self._germline_db_j = IgBLASTArgument("germline_db_j", "germline_db_J", "", True)
        self._germline_db_c = IgBLASTArgument("c_region_db", "c_region_db", "", True)
        self._aux_path = IgBLASTArgument("aux_path", "auxiliary_data", "", True)

        # Igdata is not an official blast argument, it is an enviroment
        self._igdata = Path(".")
        self.temp_dir = tmp_dir or Path(".")

    def _get_version(self) -> semantic_version.Version:
        """Private method to parse igblast -version and get semantic_version

        Returns
        -------
        semantic_version.Version
            the igblast version
        """
        try:
            process = subprocess.run([self.executable, "-version"], capture_output=True)
            stdout = process.stdout.decode("utf-8")
        except (FileNotFoundError, subprocess.CalledProcessError) as e:
            raise BadIgBLASTExe(self.executable, f"igblastn -version failed: {e}")
        version = stdout.split("\n")[0].split(":")[-1].strip()
        try:
            version = semantic_version.Version(version)
        except ValueError:
            raise BadIgBLASTExe(self.executable, f"semantic version can't parse {stdout}.")
        return version

    @property
    def version(self) -> semantic_version.Version:
        return self._version

    @property
    def executable(self) -> Path:
        """The igblastn executable

        Returns
        -------
        Path
            igblastn path
        """
        return self._executable

    @executable.setter
    def executable(self, exe: Path) -> None:
        self._executable = exe

    @property
    def temp_dir(self) -> Path:
        """The path to the tempdata directory for spliting igblast

        Returns
        -------
        Path
           A valid temporary directory path
        """
        return self._temp_dir

    @temp_dir.setter
    def temp_dir(self, data: Union[Path, str]) -> None:
        if isinstance(data, str):
            data = Path(data)
        self._temp_dir = data.absolute()
        if not os.access(self._temp_dir, os.W_OK):
            raise IOError(self._temp_dir, "Unable to write to temp dir")

    @property
    def igdata(self) -> Path:
        """The path to IGDATA which contains the internal_data needed to make a recomdination

        Returns
        -------
        Path
           A valid IGDATA path
        """
        return Path(self._igdata)

    @igdata.setter
    def igdata(self, data: Union[Path, str]) -> None:
        if isinstance(data, str):
            data = Path(data)
        if not data.exists() or not data.is_dir():
            raise BadIgDATA(data)
        self._igdata = Path(data.absolute())

    @property
    def min_d_match(self) -> ArgumentType:
        """Required minimal consecutive nucleotide base matches for D genes

        Returns
        -------
        IgBLASTArgument
        """
        return self._min_d_match

    @min_d_match.setter
    def min_d_match(self, d: int) -> None:
        """Required minimal consecutive nucleotide base matches for D genes

        Parameters
        ----------
        d : int
            The number of matches

        Raises
        ------
        BadIgBLASTArgument
            Must be >= 5
        """
        if d < 5:
            raise BadIgBLASTArgument(d, ">= 5")
        self._min_d_match = IgBLASTArgument("min_d_match", "min_D_match", d, False)

    @property
    def num_v(self) -> ArgumentType:
        """
           Number of Germline sequences to show alignments for

        Returns
        -------
        IgBLASTArgument
        """
        return self._num_v

    @num_v.setter
    def num_v(self, v: int) -> None:
        self._num_v = IgBLASTArgument("num_v", "num_alignments_V", v, False)

    @property
    def num_d(self) -> ArgumentType:
        """
           Number of Germline sequences to show alignments for D gene

        Returns
        -------
        IgBLASTArgument
        """
        return self._num_d

    @num_d.setter
    def num_d(self, d: int) -> None:
        self._num_d = IgBLASTArgument("num_d", "num_alignments_D", d, False)

    @property
    def num_j(self) -> ArgumentType:
        """
           Number of Germline sequences to show alignments for J gene

        Returns
        -------
        IgBLASTArgument
        """
        return self._num_j

    @num_j.setter
    def num_j(self, j: int) -> None:
        self._num_j = IgBLASTArgument("num_j", "num_alignments_J", j, False)

    @property
    def organism(self) -> ArgumentType:
        """The organism for your query sequence.

        Returns
        -------
        IgBLASTArgument
        """
        return self._organism

    @organism.setter
    def organism(self, o: str) -> None:
        """Organism

        Parameters
        ----------
        o : str
            an organism string
        """
        # I don't want to hardcode in the organisms here.
        # I will handle that logic at a higher level,
        # this is because blast has no preset organims and it's all about the v,d,j blast paths which are set dynamically
        self._organism = IgBLASTArgument("organism", "organism", o, True)

    @property
    def outfmt(self) -> ArgumentType:
        """alignment view options:
            3 = Flat query-anchored, show identities,
            4 = Flat query-anchored, no identities,
            7 = Tabular with comment lines
            19 = Rearrangement summary report (AIRR format)

        Returns
        -------
        IgBLASTArgument
        """
        return self._outfmt

    @outfmt.setter
    def outfmt(self, fmt: int) -> None:
        """alignment view options:
            3 = Flat query-anchored, show identities,
            4 = Flat query-anchored, no identities,
            7 = Tabular with comment lines
            19 = Rearrangement summary report (AIRR format)

        Exceptions
        ----------
        BadIgBLASTArgument
            Only 19 is supported for now
        """
        # only accept 19 for now
        if fmt != 19:
            raise BadIgBLASTArgument(fmt, 19)
        self._outfmt = IgBLASTArgument("outfmt", "outfmt", fmt, True)

    @property
    def receptor(self) -> ArgumentType:
        """Specify Ig or T cell receptor sequence

        Returns
        -------
        IgBLASTArgument
        """
        return self._receptor

    @receptor.setter
    def receptor(self, r: str) -> None:
        """Specify Ig or T cell receptor sequence

        Exceptions
        ----------
        BadIgBLASTArgument
            Only 'Ig' or 'TCR' are supported
        """
        if r not in ["Ig", "TCR"]:
            raise BadIgBLASTArgument(r, ["Ig", "TCR"])
        self._receptor = IgBLASTArgument("receptor", "ig_seqtype", r, True)

    @property
    def nomenclature(self) -> ArgumentType:
        """Domain system to be used for segment annotation

        Returns
        -------
        IgBLASTArgument
        """
        return self._nomenclature

    @nomenclature.setter
    def nomenclature(self, system: str) -> None:
        """Domain system to be used for segment annotation

        Exceptions
        ----------
        BadIgBLASTArgument
            Only 'imgt' or 'kabat' are supported
        """
        if system.lower() not in ["imgt", "kabat"]:
            raise BadIgBLASTArgument(system, "['imgt','kabat']")
        self._nomenclature = IgBLASTArgument("nomenclature", "domain_system", system, True)

    @property
    def aux_path(self) -> ArgumentType:
        """Auxilary data path. This is needed to lookup the J genes and tell them when the CDR3 stops.

        Returns
        -------
        IgBLASTArgument
        """
        return self._aux_path

    @aux_path.setter
    def aux_path(self, aux_path: Path | str) -> None:
        if isinstance(aux_path, str):
            aux_path = Path(aux_path)
        if not aux_path.exists():
            raise BadIgBLASTArgument(aux_path, "valid path to Auxilary database")
        self._aux_path = IgBLASTArgument("aux_path", "auxiliary_data", aux_path.absolute(), True)

    @property
    def germline_db_v(self) -> ArgumentType:
        """Path to V gene database prefix

        Returns
        -------
        IgBLASTArgument
        """
        return self._germline_db_v

    @germline_db_v.setter
    def germline_db_v(self, path: str | Path) -> None:
        """Path to V gene database prefix

        Exceptions
        ----------
        BadIgBLASTArgument
            Only valid path to V Database is supported
        """
        abs_path = ensure_prefix_to(path)
        if not abs_path:
            raise BadIgBLASTArgument(path, "Valid path to V Database")
        self._germline_db_v = IgBLASTArgument("germline_db_v", "germline_db_V", path, True)

    @property
    def germline_db_d(self) -> ArgumentType:
        """Path to D gene database prefix

        Returns
        -------
        IgBLASTArgument
        """
        return self._germline_db_d

    @germline_db_d.setter
    def germline_db_d(self, path: str | Path) -> None:
        """Path to D gene database prefix"""
        abs_path = ensure_prefix_to(path)
        if not abs_path:
            warnings.warn(f"{path} is not found, No D gene segment", UserWarning)
            # raise BadIgBLASTArgument(path, "Valid path to D Database")
            self._germline_db_d = IgBLASTArgument("germline_db_d", "germline_db_D", "", False)
        else:
            self._germline_db_d = IgBLASTArgument("germline_db_d", "germline_db_D", path, True)

    @property
    def germline_db_c(self) -> ArgumentType:
        """Path to C gene database prefix

        Returns
        -------
        IgBLASTArgument
        """
        return self._germline_db_c

    @germline_db_c.setter
    def germline_db_c(self, path: str | Path) -> None:
        """Path to C gene database prefix"""
        abs_path = ensure_prefix_to(path)
        if not abs_path:
            warnings.warn(f"{path} is not found, No C gene segment", UserWarning)
            self._germline_db_c = IgBLASTArgument("c_region_db", "c_region_db", "", False)
        else:
            self._germline_db_c = IgBLASTArgument("c_region_db", "c_region_db", path, True)

    @property
    def germline_db_j(self) -> ArgumentType:
        """Path to J gene database prefix

        Returns
        -------
        IgBLASTArgument
        """
        return self._germline_db_j

    @germline_db_j.setter
    def germline_db_j(self, path: str | Path) -> None:
        """Path to J gene database prefix

        Exceptions
        ----------
        BadIgBLASTArgument
            Only valid path to J Database is supported
        """
        abs_path = ensure_prefix_to(path)
        if not abs_path:
            raise BadIgBLASTArgument(path, "Valid path to J Database")
        self._germline_db_j = IgBLASTArgument("germline_db_j", "germline_db_J", path, True)

    @property
    def word_size(self) -> ArgumentType:
        """Word size for wordfinder algorithm (length of best perfect match)

        Returns
        -------
        IgBLASTArugment
        """
        return self._word_size

    @word_size.setter
    def word_size(self, word_size: int) -> None:
        """Word size for wordfinder algorithm (length of best perfect match)

        Exceptions
        ----------
        BadIgBLASTArgument
            Only integer greater than 4 is supported
        """
        if word_size < 4:
            raise BadIgBLASTArgument(word_size, ">=4")
        self._word_size = IgBLASTArgument("word_size", "word_size", word_size, False)

    @property
    def gap_open(self) -> ArgumentType:
        """Cost to open a gap

        Returns
        -------
        IgBLASTArgument
        """
        return self._gap_open

    @gap_open.setter
    def gap_open(self, go: int) -> None:
        """Cost to open a gap

        Exceptions
        ----------
        BadIgBLASTArgument
            Only integer greater than 0 is supported
        """
        if go < 0:
            raise BadIgBLASTArgument(go, ">=0")
        self._gap_open = IgBLASTArgument("gap_open", "gapopen", go, False)

    @property
    def gap_extend(self) -> ArgumentType:
        """Cost to extend a gap

        Returns
        -------
        IgBLASTArgument
        """
        return self._gap_extend

    @gap_extend.setter
    def gap_extend(self, ge: int) -> None:
        """Cost to extend a gap

        Exceptions
        ----------
        BadIgBLASTArgument
            Only integer greater than 0 is supported
        """
        if ge < 0:
            raise BadIgBLASTArgument(ge, ">=0")
        self._gap_extend = IgBLASTArgument("gap_open", "gapextend", ge, False)

    @property
    def num_threads(self) -> ArgumentType:
        """
            Number of threads (CPUs) to use in the BLAST search

        Returns
        -------
        IgBLASTArgument
        """
        return self._num_threads

    @num_threads.setter
    def num_threads(self, num_threads: int) -> None:
        # Some architectures have more threads than number of cores. Letting the users decide the number of threads.
        self._num_threads = IgBLASTArgument("number_threds", "num_threads", num_threads, False)

    @property
    def extend_5(self) -> ArgumentType:
        """Extend V gene alignment at 5' end

        Returns
        -------
        IgBLASTArgument
        """
        return self._extend_5

    @extend_5.setter
    def extend_5(self, extend_5: bool) -> None:
        self._extend_5 = IgBLASTArgument("extend_5", "extend_align5end", extend_5, False)

    @property
    def extend_3(self) -> ArgumentType:
        """Extend V gene alignment at 3' end

        Returns
        -------
        IgBLASTArgument
        """
        return self._extend_3

    @extend_3.setter
    def extend_3(self, extend_3: bool) -> None:
        self._extend_3 = IgBLASTArgument("extend_3", "extend_align3end", extend_3, False)

    @property
    def allow_vdj_overlap(self) -> Any:
        """Allow the VDJ overlap

        This option is active only when D_penalty
        and J_penalty are set to -4 and -3, respectively

        Returns
        -------
        IgBLASTArgument
        """
        return self._allow_vdj_overlap  # type: ignore[has-type]

    @allow_vdj_overlap.setter
    def allow_vdj_overlap(self, allow: bool) -> None:
        """Allow the VDJ overlap

        This option is active only when D_penalty
        and J_penalty are set to -4 and -3, respectively
        """
        j_penalty: IgBLASTArgument = self.j_penalty  # type: ignore[assignment]
        d_penalty: IgBLASTArgument = self.d_penalty  # type: ignore[assignment]
        if j_penalty.value != -3 and d_penalty.value != -4 and allow:
            warnings.warn(
                f"Allows vdj overlap set but j penalty and d penalty need to be -3 and -4, now are {self.j_penalty}, {self.d_penalty}",
                UserWarning,
            )
        self._allow_vdj_overlap = IgBLASTArgument("allow_vdj_overlap", "allow_vdj_overlap", allow, False)

    @property
    def d_penalty(self) -> ArgumentType:
        """What is the  D gene panalty

        Returns
        -------
        IgBLASTArgument
        """
        return self._d_penalty

    @d_penalty.setter
    def d_penalty(self, penalty: int) -> None:
        """What is the  D gene panalty

        Exceptions
        ----------
        BadIgBLASTArgument
            Only integer less than 0 and greater than -5 is supported
        """
        if not -5 < penalty < 1:
            raise BadIgBLASTArgument(penalty, "must be less than 0 and greater than -5")
        self._d_penalty = IgBLASTArgument("d_penalty", "D_penalty", penalty, True)

    @property
    def j_penalty(self) -> ArgumentType:
        """What is the  J gene panalty

        Returns
        -------
        IgBLASTArgument
        """
        return self._j_penalty

    @j_penalty.setter
    def j_penalty(self, penalty: int) -> None:
        """What is the  J gene panalty

        Exceptions
        ----------
        BadIgBLASTArgument
            Only integer less than 0 and greater than -4 is supported
        """
        if not -4 < penalty < 1:
            raise BadIgBLASTArgument(penalty, "must be less than 0 and greater than -4")
        self._j_penalty = IgBLASTArgument("j_penalty", "J_penalty", penalty, True)

    @property
    def v_penalty(self) -> ArgumentType:
        """What is the  v gene panalty

        Returns
        -------
        IgBLASTArgument
        """
        return self._v_penalty

    @v_penalty.setter
    def v_penalty(self, penalty: int) -> None:
        if not -5 < penalty < 1:
            raise BadIgBLASTArgument(penalty, "must be less than 0 and greater than -5")
        self._v_penalty = IgBLASTArgument("v_penalty", "V_penalty", penalty, True)

    @property
    def arguments(self) -> List[IgBLASTArgument]:
        """return a list of IgBLASTArugments

        Returns
        -------
        List[IgBLASTArguments]
        """
        # lots of type ignores since these are IgBLASTArguments set in the setter, but are read from the property
        return [
            self.min_d_match,  # type: ignore
            self.num_v,  # type: ignore
            self.num_j,  # type: ignore
            self.num_d,  # type: ignore
            self.organism,  # type: ignore
            self.receptor,  # type: ignore
            self.germline_db_v,  # type: ignore
            self.germline_db_d,  # type: ignore
            self.germline_db_j,  # type: ignore
            self.germline_db_c,  # type: ignore
            self.aux_path,  # type: ignore
            self.outfmt,  # type: ignore
            self.nomenclature,  # type: ignore
            self.word_size,  # type: ignore
            self.gap_open,  # type: ignore
            self.gap_extend,  # type: ignore
            self.j_penalty,  # type: ignore
            self.v_penalty,  # type: ignore
            self.d_penalty,  # type: ignore
            self.num_threads,  # type: ignore
            self.extend_5,  # type: ignore
            self.extend_3,  # type: ignore
            self.allow_vdj_overlap,
        ]

    @property
    def cmd(self) -> List[str]:
        """Return the blast cmd that will be run by subprocess"""
        _cmd = [str(self.executable)]
        for blast_arg in self.arguments:
            kv = blast_arg.get_formatted_blast_arg()  # can return non if we already set it twice
            if kv:  # only set on boolean if they are true
                _cmd += kv
        return _cmd

    def pre_check(self) -> None:
        """Ensures we have set everything right

        Raises
        ------
        MissingIgBLASTArg
            We have set the IGDATA field
        BadIgBLASTExe
            Correct IGblast executable
        BadIgDATA
            If any of the fields are not set properly
        """
        # TODO: move to pydanic model
        # Ensure required arguments werer set
        for blast_arg in self.arguments:
            if blast_arg.required and not (blast_arg.value):
                raise MissingIgBLASTArgument(f"Missing Blast argument. Need to set IgBLASTN.{blast_arg.name}")

    # Run methods
    def run_file(self, file: Union[Path, str]) -> pd.DataFrame:
        """Run IgBlast on a file

        Parameters
        ----------
        file : Path
            the fasta file path

        Returns
        -------
        pd.DataFrame
            A dataframe with the IgBLAST results

        Raises
        ------
        IgBLASTRunTimeError
           for any given runtime error for igblastn
        """

        # because igblast uses IGDATA as the internal file structure, we should pass the enviroment to the subprocess
        local_env = os.environ.copy()
        local_env["IGDATA"] = str(self.igdata)

        # take the cmd and finally add the query file
        cmd = self.cmd
        cmd += ["-query", str(file)]

        # run a precheck to make sure everything passed was working
        self.pre_check()

        # while we can certainly do this as an output stream on stdout,
        # It's probably best to take advantage of IGblast output and tempfile
        with tempfile.NamedTemporaryFile(dir=self.temp_dir, suffix="_igblast.tsv") as tmpfile:
            cmd += ["-out", tmpfile.name]
            process = subprocess.run(cmd, env=local_env, capture_output=True)
            if process.stderr:
                raise IgBLASTRunTimeError(process.stderr)
            # we read the dataframe from the tempfile, it should always be in .TSV.
            # We can also cast it to IGBLAST_AIRR dtypes to save memory
            df = pd.read_csv(tmpfile.name, sep="\t", dtype=IGBLAST_AIRR)  # type: ignore

        df["v_identity"] = df["v_identity"] / 100
        df["d_identity"] = df["d_identity"] / 100
        df["j_identity"] = df["j_identity"] / 100
        return df

    def __repr__(self) -> str:
        return "IgBLAST: env IGDATA={} {}".format(str(self.igdata), " ".join(self.cmd))

    def __str__(self) -> str:
        return self.__repr__()


#!/user/bin/python
# -*- coding:utf-8 -*-
# Scopus-test.py
# 获取大学排名信息

"""
http://www.zuihaodaxue.cn/zuihaodaxuepaiming2018.html

功能描述：
    输入：大学排名URL
    输出：大学排名信息的输出（排名，大学名称，地区，总分）
技术路径：requests-bs4
定向爬虫：仅对输入URL进爬取，不扩展爬取其他URL。

步骤1：从网络上获取大学排名网页内容
    getHTMLText()
步骤2：提取网页内容中信息到合适的数据结构
    fillUnivList()
步骤3：利用数据结构展示并输出结果
    printUnivList()

format 方法回顾：【:<><><>,<>,<>】

"""

import requests
import bs4

# 步骤1：从网络上获取大学排名网页内容
def getHTMLText(url):
    try:
        r = requests.get(url,timeout = 30)
        r.raise_for_status()
        r.encoding = r.apparent_encoding
        return r.text
    except:
        return ""

# 步骤2：提取网页内容中信息到合适的数据结构
def fillUnivList(ulist,html):
    soup = bs4.BeautifulSoup(html, "html.parser")
    for tr in soup.find("tbody").children:
        if isinstance(tr, bs4.element.Tag):
            td_trs = []
            for td in tr:
                td_trs.append(td)
            ulist.append([td_trs[0].string, td_trs[1].div.string, td_trs[2].string, td_trs[3].string])

# 步骤3：利用数据结构展示并输出结果
def printUnivList(ulist,num):
    '''
    中英文混合输出时，采用中文字符的空格填充 chr(12288)
    :param ulist:
    :param num:
    :return:
    '''

    #生成输出模板
    tplt = "{0:^5}\t{1:{4}^10}\t{2:<6}\t{3:<10}"
    print(tplt.format("排名", "学校名称", "地区", "总分", chr(12288)))
    for i in range(num):
        u = ulist[i]
        print(tplt.format(u[0], u[1], u[2], u[3], chr(12288)))

def main():
    uinfo = []
    url = "http://www.zuihaodaxue.cn/zuihaodaxuepaiming2018.html"
    html = getHTMLText(url)
    fillUnivList(uinfo, html)
    printUnivList(uinfo, 50)    # 打印20个

# 执行
main()


import streamlit as st
from multipage import MultiApp
from apps import pandas_profiling_uk_houseprice, web_app_uk_houseprice

app = MultiApp()

st.set_page_config(
    page_title="UK House Price",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
# UK House Price App
With the launch of Streamlit, developing a dashboard for your machine learning solution has been made incredibly easy.

Streamlit is an open source app framework specifically designed for ML engineers working with Python.
It allows you to create a stunning looking application with only a few lines of code.

I want to take this opportunity to demonstrate the apps you can build using Streamlit.
""")

# Add all your application here
app.add_app("Pandas Profiling UK Houseprice", pandas_profiling_uk_houseprice.app)
app.add_app("Trend UK Houseprice", web_app_uk_houseprice.app)
# app.add_app("Statistical Model", model.app)
# The main app
app.run()


import logging
import math

from gql import Client, gql
from gql.transport.aiohttp import AIOHTTPTransport

_LOGGER = logging.getLogger(__name__)

class OctopusEnergyGraphQLClient:

  def __init__(self, api_key: str):
    if (api_key == None):
      raise Exception('API KEY is not set')

    self._api_key = api_key
    self._base_url = "https://api.octopus.energy/v1/graphql/"
    self._login_attempt = 0
    self._session = None

  async def async_get_accounts(self):
    """Gets the accounts for the given API key"""
    return await self.__async_execute_with_session(self.__async_get_accounts)

  async def async_get_combined_state(self, account_id: str):
    """Gets the state for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_get_combined_state(session, account_id))

  async def async_get_charge_preferences(self, account_id: str):
    """Gets the charging preferences for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_get_charge_preferences(session, account_id))

  async def async_set_charge_preferences(self, account_id: str, readyByHoursAfterMidnight: float, targetSocPercent: int):
    """Sets the charging preferences for the given account"""
    return await self.__async_execute_with_session(
      lambda session: self.__async_set_charge_preferences(session,
        account_id,
        readyByHoursAfterMidnight,
        targetSocPercent))

  async def async_trigger_boost_charge(self, account_id: str):
    """Triggers a boost charge for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_trigger_boost_charge(session, account_id))

  async def async_cancel_boost_charge(self, account_id: str):
    """Cancels the boost charge for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_cancel_boost_charge(session, account_id))

  async def async_suspend_smart_charging(self, account_id: str):
    """Suspends smart charging for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_suspend_smart_charging(session, account_id))

  async def async_resume_smart_charging(self, account_id: str):
    """Resumes smart charging for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_resume_smart_charging(session, account_id))

  async def async_get_device_info(self, account_id: str):
    """Gets the device info for the given account"""
    return await self.__async_execute_with_session(lambda session: self.__async_get_device_info(session, account_id))

  async def __async_get_token(self):
    """Gets a new token from the API"""
    transport = AIOHTTPTransport(url=self._base_url)

    # Using `async with` on the client will start a connection on the transport
    # and provide a `session` variable to execute queries on this connection
    async with Client(
      transport=transport,
      fetch_schema_from_transport=True,
    ) as session:

      # Execute single query
      query = gql(
      '''
        mutation krakenTokenAuthentication($apiKey: String!) {
          obtainKrakenToken(input: { APIKey: $apiKey })
          {
            token
          }
        }
      ''')

      params = {"apiKey": self._api_key}

      result = await session.execute(query, variable_values=params, operation_name="krakenTokenAuthentication")
      token = result['obtainKrakenToken']['token']
      return token

  async def __async_get_session(self, reset = False):
    """Creates a new session, use with `async with`"""
    if (reset):
      self._session = None

    if (self._session != None):
      return self._session

    token = await self.__async_get_token()
    headers = {"Authorization": token}
    transport = AIOHTTPTransport(url=self._base_url, headers=headers)

    # Using `async with` on the client will start a connection on the transport
    # and provide a `session` variable to execute queries on this connection
    session = Client(
      transport=transport,
      fetch_schema_from_transport=True,
    )
    self._session = session
    return session

  async def __async_execute_with_session(self, func):
    """Executes the given function with a session, auto retrying with a new session if it fails."""
    try:
      async with await self.__async_get_session() as session:
        return await func(session)
    except Exception as e:
      try:
        async with await self.__async_get_session(reset = True) as session:
          return await func(session)
      except Exception as e:
        raise e

  async def __async_set_charge_preferences(self, session, account_id: str, readyByHoursAfterMidnight: float, targetSocPercent: int):
    """Sets the charging preferences for the given account"""
    # round up to nearest 5
    targetSocPercent = 5 * math.ceil(round(targetSocPercent) / 5)
    # round up to nearest 0.5
    readyByHoursAfterMidnight = 0.5 * round(readyByHoursAfterMidnight / 0.5)

    if readyByHoursAfterMidnight < 4 or readyByHoursAfterMidnight > 11:
        raise ValueError("Target time must be between 4AM and 11AM")
    if targetSocPercent < 10 or targetSocPercent > 100:
        raise ValueError("Target SOC percent must be between 10 and 100")

    # get hours
    readyByHoursAfterMidnightHours = int(readyByHoursAfterMidnight)
    # get mins - rounded to nearest 30 mins
    readyByHoursAfterMidnightMinutes = round(60 * (readyByHoursAfterMidnight % 1))

    targetTime = f"{readyByHoursAfterMidnightHours:02}:{readyByHoursAfterMidnightMinutes:02}"

    # Execute single query
    query = gql(
    '''
      mutation setVehicleChargePreferences($accountNumber: String!, $targetTime: String!, $targetSocPercent: Int!) {
        setVehicleChargePreferences(input: { accountNumber: $accountNumber, weekdayTargetTime: $targetTime, weekendTargetTime: $targetTime, weekdayTargetSoc: $targetSocPercent, weekendTargetSoc: $targetSocPercent }) {
          krakenflexDevice {
            krakenflexDeviceId
          }
        }
      }
    ''')

    params = {"accountNumber": account_id, "targetTime": targetTime, "targetSocPercent": targetSocPercent}
    result = await session.execute(query, variable_values=params, operation_name="setVehicleChargePreferences")
    return result['setVehicleChargePreferences']

  async def __async_trigger_boost_charge(self, session, account_id: str):
    """Triggers a boost charge for the given account"""
    # Execute single query
    query = gql(
    '''
      mutation triggerBoostCharge($accountNumber: String!) {
        triggerBoostCharge(input: { accountNumber: $accountNumber }) {
          krakenflexDevice {
            krakenflexDeviceId
          }
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="triggerBoostCharge")
    return result['triggerBoostCharge']

  async def __async_cancel_boost_charge(self, session, account_id: str):
    """Cancels any boost charge currently in progress for the given account"""
    # Execute single query
    query = gql(
    '''
      mutation deleteBoostCharge($accountNumber: String!) {
        deleteBoostCharge(input: { accountNumber: $accountNumber }) {
          krakenflexDevice {
            krakenflexDeviceId
          }
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="deleteBoostCharge")
    return result['deleteBoostCharge']

  async def __async_get_accounts(session):
    # Execute single query
    query = gql(
    '''
        query viewer {
          viewer {
            fullName
            accounts {
              number
            }
          }
        }
    ''')

    params = {}
    result = await session.execute(query, variable_values=params, operation_name="viewer")

    return list(map(lambda x: x['number'], result['viewer']['accounts']))

  async def __async_get_charge_preferences(self, session, account_id: str):
    """Gets the charging preferences for the given account"""
    # Execute single query
    query = gql(
    '''
      query vehicleChargingPreferences($accountNumber: String!) {
        vehicleChargingPreferences(accountNumber: $accountNumber) {
          weekdayTargetTime,
          weekdayTargetSoc,
          weekendTargetTime,
          weekendTargetSoc
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="vehicleChargingPreferences")
    return result['vehicleChargingPreferences']

  async def __async_get_device_info(self, session, account_id: str):
    """Get the user's device info (e.g. vehicle make, battery size, state etc.)"""
    # Execute single query
    query = gql(
    '''
      query registeredKrakenflexDevice($accountNumber: String!) {
        registeredKrakenflexDevice(accountNumber: $accountNumber) {
          krakenflexDeviceId
          provider
          vehicleMake
          vehicleModel
          vehicleBatterySizeInKwh
          chargePointMake
          chargePointModel
          chargePointPowerInKw
          status
          suspended
          hasToken
          createdAt
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="registeredKrakenflexDevice")
    return result['registeredKrakenflexDevice']

  async def __async_get_combined_state(self, session, account_id: str):
    """Get the user's account state"""
    query = gql(
    '''
        query getCombinedData($accountNumber: String!) {
            vehicleChargingPreferences(accountNumber: $accountNumber) {
                weekdayTargetTime,
                weekdayTargetSoc,
                weekendTargetTime,
                weekendTargetSoc
            }
            registeredKrakenflexDevice(accountNumber: $accountNumber) {
                krakenflexDeviceId
                provider
                vehicleMake
                vehicleModel
                vehicleBatterySizeInKwh
                chargePointMake
                chargePointModel
                chargePointPowerInKw
                status
                suspended
                hasToken
                createdAt
            }
            plannedDispatches(accountNumber: $accountNumber) {
                startDtUtc: startDt
                endDtUtc: endDt
                chargeKwh: delta
                meta {
                    source
                    location
                }
            }
            completedDispatches(accountNumber: $accountNumber) {
                startDtUtc: startDt
                endDtUtc: endDt
                chargeKwh: delta
                meta {
                    source
                    location
                }
            }
        }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="getCombinedData")
    return result

  async def __async_suspend_smart_charging(self, session, account_id: str):
    """Suspends smart charging for the given account"""
    # Execute single query
    query = gql(
    '''
      mutation suspendControl($accountNumber: String!) {
        suspendControl(input: { accountNumber: $accountNumber }) {
          krakenflexDevice {
            krakenflexDeviceId
          }
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="suspendControl")
    return result['suspendControl']

  async def __async_resume_smart_charging(self, session, account_id: str):
    """Resumes smart charging for the given account"""
    # Execute single query
    query = gql(
    '''
      mutation resumeControl($accountNumber: String!) {
        resumeControl(input: { accountNumber: $accountNumber }) {
          krakenflexDevice {
            krakenflexDeviceId
          }
        }
      }
    ''')

    params = {"accountNumber": account_id}
    result = await session.execute(query, variable_values=params, operation_name="resumeControl")
    return result['resumeControl']


import matplotlib.pyplot as plt
import numpy as np

n_list = [1, 2, 3, 4, 5, 6, 7, 8]
a_list = [2804, 2530, 6313, 1682, 3310, 3608, 2647, 2278]
m_list = [2158, 939, 3512, 1179, 1685, 2298, 1559, 1310]

plt.plot(n_list, a_list, color='b', marker='.', label='std::map')
plt.plot(n_list, m_list, color='r', marker='.', label='RBtree')

plt.legend(loc=2, borderaxespad=1)

plt.xlabel('Test num')
plt.ylabel('Time(μs)')
plt.title('%%time')
# plt.grid(True)
plt.savefig("time.png")
plt.show()


#!/usr/bin/env python2.5
import hildon
import gtk, gobject
import os
from album_art_spec import getCoverArtThumbFileName
from download_dialog import MussorgskyAlbumArtDownloadDialog
from utils import escape_html
from aa_selection_dialog import AlbumArtSelectionDialog, RESPONSE_CLICK

import time

EMPTY_PIXBUF = gtk.gdk.Pixbuf (gtk.gdk.COLORSPACE_RGB, False, 8, 64, 64)

import i18n
_ = i18n.language.gettext

class MussorgskyAlbumArtPanel (hildon.StackableWindow):

    def __init__ (self, album_artists):
        hildon.StackableWindow.__init__ (self)
        self.set_title (_("Album art selection"))
        self.set_border_width (12)
        self.__create_view ()
        self.downloader = None
        # Visible string, image, artist, album, painted!
        self.model = gtk.ListStore (str, gtk.gdk.Pixbuf, str, str, bool)
        for p in album_artists:
            if (not p[0]):
                continue
            t = (None, None, p[1], p[0], False)
            self.model.append (t)

        self.treeview.set_model (self.model)

    def __create_view (self):
        self.treeview = gtk.TreeView ()
        self.treeview.connect ("row-activated", self.row_activated_cb)

        artist_column = gtk.TreeViewColumn ("Artist", gtk.CellRendererText (), markup=0)
        artist_column.set_expand (True)
        self.treeview.append_column (artist_column)

        renderer = gtk.CellRendererPixbuf ()
        album_art = gtk.TreeViewColumn ("Album art", renderer, pixbuf=1)
        # This doesn't have real effect:
        album_art.set_sizing (gtk.TREE_VIEW_COLUMN_FIXED)
        album_art.set_fixed_width (64)

        album_art.set_cell_data_func (renderer, self.album_art_cell_data_cb)
        self.treeview.append_column (album_art)

        pannable_area = hildon.PannableArea ()
        pannable_area.add (self.treeview)
        self.add (pannable_area)

        # Menu
        menu = hildon.AppMenu ()
        automatic_retrieval = hildon.Button (hildon.BUTTON_STYLE_NORMAL,
                                             hildon.BUTTON_ARRANGEMENT_HORIZONTAL)
        automatic_retrieval.set_title (_("Automatic download"))
        automatic_retrieval.connect ("clicked", self.get_all_album_art)
        menu.append (automatic_retrieval)
        menu.show_all ()
        self.set_app_menu (menu)

    def album_art_cell_data_cb (self, column, cell, model, iter):
        text, pixbuf, artist, album, not_first_time = model.get (iter, 0, 1, 2, 3, 4)
        if (not_first_time):
            if (text == None):
                text = "".join (["", escape_html (album),"\n",
                                 escape_html(artist), ""])
                model.set (iter, 0, text)

            if (pixbuf == None):
                #print "Calling album art cell data cb", model.get (iter, 3)
                album_art_path = getCoverArtThumbFileName (album)
                if (os.path.exists (album_art_path)):
                    pxb = gtk.gdk.pixbuf_new_from_file_at_size (album_art_path, 64, 64)
                    model.set (iter, 1, pxb)
                else:
                    #print "Cannot find thumbnail in '%s'" % (album_art_path)
                    model.set (iter, 1, EMPTY_PIXBUF)

        else:
            model.set (iter, 4, True)

    def get_all_album_art (self, user_data):
        dialog = MussorgskyAlbumArtDownloadDialog (self)
        dialog.show_all ()
        dialog.do_the_job (self.model)

    def row_activated_cb (self, treeview, path, view_colum):
        it = treeview.get_model ().get_iter (path)
        artist = treeview.get_model ().get_value (it, 2)
        album = treeview.get_model ().get_value (it, 3)
        if (artist.find ('|') != -1):
            artist = "Various artists"

        dialog = AlbumArtSelectionDialog (self, artist, album, 5)
        dialog.show_all ()

        response = dialog.run ()
        if (response == RESPONSE_CLICK):
            (img, thumb) = dialog.get_selection ()
            if img and thumb:
                pixbuf = gtk.gdk.pixbuf_new_from_file_at_size (thumb, 64, 64)
                treeview.get_model ().set (it, 1, pixbuf)
            else:
                treeview.get_model ().set (it, 1, EMPTY_PIXBUF)
        dialog.destroy ()

if __name__ == "__main__":
    import random

    artists_albums = [("Artist %d|artist Y" % i, "Album <%d>" % i) for i in range (0, 10)]

    # Overwrite the get thumb path for testing
    def local_file (path):
        return "../thumb%d.124.jpeg" % (random.randint (0, 3))

    global getCoverArtThumbFileName
    getCoverArtThumbFileName = local_file

    window = MussorgskyAlbumArtPanel (artists_albums)
    window.connect ("destroy", gtk.main_quit )
    window.show_all ()
    gtk.main ()


from rest_framework import serializers

from .models import Comment

class CommentSerializer(serializers.ModelSerializer):
    username = serializers.CharField(source='user.username', read_only=True)

    class Meta:
        model = Comment
        fields = [
            'id',
            'modified_at',
            'body',
            'post',
            'user',
            'username'
        ]


"""
Used to visualise tracks and detections

"""
from absl import app, flags, logging
from absl.flags import FLAGS

import cv2
import math
import numpy as np
import os
import queue
import random
from tqdm import tqdm

from visualisation.image import cv_plot_bbox

def visualise(video_path, frames_dir, detections_dir, tracks_dir, stats_dir, vis_dir,
              img_snapshots_dir, vid_snapshots_dir, around='detections',
              start_buffer=100, end_buffer=50,
              display_tracks=False, display_detections=False, display_trails=False, save_static_trails=False,
              generate_image_snapshots=False, generate_video_snapshots=False, summary=True, full=False):

    video_path = os.path.normpath(video_path)
    frames_dir = os.path.normpath(frames_dir)
    detections_dir = os.path.normpath(detections_dir)
    tracks_dir = os.path.normpath(tracks_dir)
    stats_dir = os.path.normpath(stats_dir)
    vis_dir = os.path.normpath(vis_dir)
    img_snapshots_dir = os.path.normpath(img_snapshots_dir)
    vid_snapshots_dir = os.path.normpath(vid_snapshots_dir)

    # if not generate_image_snapshots:
    #     img_snapshots_dir = None
    # if not generate_video_snapshots:
    #     vid_snapshots_dir = None

    colors = dict()
    for i in range(200):
        colors[i] = (int(256 * random.random()), int(256 * random.random()), int(256 * random.random()))

    os.makedirs(vis_dir, exist_ok=True)
    os.makedirs(img_snapshots_dir, exist_ok=True)
    os.makedirs(vid_snapshots_dir, exist_ok=True)

    video_path, video_filename = os.path.split(video_path)
    txt_filename = video_filename[:-4]+'.txt'

    with open(os.path.join(stats_dir, txt_filename), 'r') as f:
        video_id, width, height, length = f.read().rstrip().split(',')

    # load detections
    with open(os.path.join(detections_dir, txt_filename), 'r') as f:
        detections = [line.rstrip().split(',') for line in f.readlines()]

    width = int(width)
    height = int(height)
    length = int(length)

    mult = True
    detections_ = dict()
    for d in detections:
        if mult:
            d_ = [int(d[1]), float(d[2]),
                  float(d[3])*width, float(d[4])*height,
                  float(d[5])*width, float(d[6])*height]
        else:
            d_ = [int(d[1]), float(d[2]),
                  float(d[3]), float(d[4]),
                  float(d[5]), float(d[6])]

        if int(d[0]) in d:
            detections_[int(d[0])].append(d_)
        else:
            detections_[int(d[0])] = [d_]
    detections = detections_

    # load tracks
    with open(os.path.join(tracks_dir, txt_filename), 'r') as f:
        tracks = [line.rstrip().split(',') for line in f.readlines()]

    tracks_ = dict()
    for t in tracks:
        if mult:
            t_ = [int(t[1]), float(t[2]),
                  float(t[3])*width, float(t[4])*height,
                  float(t[5])*width, float(t[6])*height]
        else:
            t_ = [int(t[1]), float(t[2]),
                  float(t[3]), float(t[4]),
                  float(t[5]), float(t[6])]

        if int(t[0]) in tracks_:
            tracks_[int(t[0])].append(t_)
        else:
            tracks_[int(t[0])] = [t_]

    tracks = tracks_

    # load video
    capture = cv2.VideoCapture(os.path.join(video_path, video_filename))
    #
    # # Get the total number of frames
    # total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    # # Might be a problem if video has no frames
    # if total < 1:
    #     logging.error("Check your opencv + ffmpeg installation, can't read videos!!!\n"
    #                   "\nYou may need to install open cv by source not pip")
    #     return None
    #
    # assert total == length-1 or total == length or total == length+1

    summary_out_video = None
    full_out_video = None
    if summary:
        summary_out_video = cv2.VideoWriter("{}_summary.mp4".format(os.path.join(vis_dir, video_filename[:-4])),
                                            cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 25, (width, height))
    if full:
        full_out_video = cv2.VideoWriter("{}_full.mp4".format(os.path.join(vis_dir, video_filename[:-4])),
                                         cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 25, (width, height))

    track_trails = queue.Queue(maxsize=50)
    if save_static_trails:
        static_track_trails = {}
        avg_frame = np.zeros((1, height, width, 3), dtype=int)

    img_track_snapshots = {}
    vid_track_snapshots = {}

    since = 0
    out_count = 0
    capture.set(1, 0)
    while_safety = 0
    for current in tqdm(range(length), desc="Visualising video: {}".format(video_filename)):

        while True:
            while_safety += 1
            flag, frame = capture.read()
            # if flag != 0 and frame is not None:
            if frame is not None:
                while_safety = 0
                break
            if while_safety > 1000:
                break

        if frame is None:  # should only occur at the end of a video
            break

        if full_out_video is not None or summary_out_video is not None:
            v_height, v_width, _ = frame.shape
            frame[-50:, -250:, :] = (0, 0, 0)
            cv2.putText(frame, '{}'.format(current), (v_width-240, v_height-12), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2)

            forward_buffer = False
            if around == 'tracks':
                if current not in tracks:
                    since += 1
                else:
                    since = 0
                for check_forward in range(current, current + start_buffer):
                    if check_forward in tracks:
                        forward_buffer = True
                        break

            elif around == 'detections':
                if current not in detections:
                    since += 1
                else:
                    since = 0
                for check_forward in range(current, current + start_buffer):
                    if check_forward in detections:
                        forward_buffer = True
                        break
            else:
                ValueError()

            if summary and not forward_buffer and since > end_buffer:
                continue  # we don't want to save out this frame

            if save_static_trails and avg_frame.shape[0] < 250 and current % 50 == 0 and frame.shape == avg_frame.shape[1:]:
                avg_frame = np.vstack((avg_frame, np.expand_dims(frame, 0)))

            out_frame = frame.copy()

            if display_tracks:
                if current in tracks:
                    out_frame = cv_plot_bbox(out_path=None,
                                             img=out_frame,
                                             bboxes=[t[2:] for t in tracks[current]],
                                             scores=[t[1] for t in tracks[current]],  # todo should do in tracking code
                                             labels=[t[0] for t in tracks[current]],
                                             thresh=0,
                                             colors=colors,
                                             class_names=[])

                    if display_trails:
                        track_trails_frame = {}
                        for t in tracks[current]:
                            # make trails per track
                            track_trails_frame[t[0]] = (int(t[2] + ((t[4] - t[2]) / 2)), int(t[5]))

                        # put in the queue that exists over frames
                        if track_trails.full():
                            track_trails.get()
                        track_trails.put(track_trails_frame)

                        # draw the trail as dots that fade with time
                        for i, trails in enumerate(list(track_trails.queue)):
                            alpha = math.pow(i / len(list(track_trails.queue)), 2)
                            overlay = out_frame.copy()
                            for tid, dot in trails.items():
                                cv2.circle(overlay, dot, 2, colors[tid], -1)
                            out_frame = cv2.addWeighted(overlay, alpha, out_frame, 1 - alpha, 0)

                    if save_static_trails:
                        for t in tracks[current]:
                            # make trails per track
                            if t[0] in static_track_trails:
                                static_track_trails[t[0]].append((int(t[2] + ((t[4] - t[2]) / 2)), int(t[5])))
                            else:
                                static_track_trails[t[0]] = [(int(t[2] + ((t[4] - t[2]) / 2)), int(t[5]))]

            if display_detections:
                if current in detections:
                    out_frame = cv_plot_bbox(out_path=None,
                                             img=out_frame,
                                             bboxes=[d[2:] for d in detections[current]],
                                             scores=[d[1] for d in detections[current]],
                                             labels=[d[0] for d in detections[current]],
                                             thresh=0,
                                             colors={0: (1, 255, 1)},
                                             class_names=['cyclist'])

            # write out the main frame
            if summary_out_video is not None:
                summary_out_video.write(out_frame)
            if full_out_video is not None:
                full_out_video.write(out_frame)
            out_count += 1

        if generate_image_snapshots and current in tracks:
            for t in tracks[current]:
                # save the part of the frame containing the cyclist
                x1, y1, x2, y2 = int(t[2]), int(t[3]), int(t[4]), int(t[5])
                if t[0] not in img_track_snapshots:
                    img_track_snapshots[t[0]] = frame[y1:y2, x1:x2, :]
                # else:  # never replace first one
                #     h, w, _ = img_track_snapshots[t[0]].shape
                #     # replace if has a bigger area, we are assuming the bigger the better
                #     if (x2 - x1) * (y2 - y1) > w * h and x1 > 0 and y1 > 0 and x2 < width and y2 < height:
                #         img_track_snapshots[t[0]] = frame[y1:y2, x1:x2, :]

        if generate_video_snapshots and current in tracks:
            # if the track wasn't found it has died, close it's associated clip
            current_tracks = [t[0] for t in tracks[current]]
            del_tids = []
            for k, vid in vid_track_snapshots.items():
                if k not in current_tracks:
                    del_tids.append(k)
                    vid.release()

            for k in del_tids:
                del vid_track_snapshots[k]

            # write out frames and add new clip if needed
            for t in tracks[current]:
                tid = t[0]
                # open a new clip for this track
                if tid not in vid_track_snapshots:
                    vid_track_snapshots[tid] = cv2.VideoWriter(
                        "%s_%d.mp4" % (os.path.join(vid_snapshots_dir, video_filename[:-4]), tid),  # open_vids),
                        cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 25, (width, height))

                # get a clean frame
                clean_frame = frame.copy()

                # write out a frame
                clean_frame = cv_plot_bbox(out_path=None,
                                           img=clean_frame,
                                           bboxes=[t[2:]],
                                           scores=[t[1]],
                                           labels=[tid],
                                           thresh=0,
                                           colors=colors,
                                           class_names=[str(tid)])

                vid_track_snapshots[tid].write(clean_frame)

    if full_out_video is not None:
        full_out_video.release()
    if summary_out_video is not None:
        summary_out_video.release()
    if capture is not None:
        capture.release()

    if out_count > 0:
        logging.info("\n\nOriginal video length: {}\nNew video length: {} ({}% of original)".format(
            length, out_count, int(100*float(out_count)/length)))
    else:
        logging.info("\n\nNo output video saved")

    # write out the snapshot images
    if generate_image_snapshots:
        for tid, img in img_track_snapshots.items():
            out_path = os.path.join(img_snapshots_dir, "{}_{:03d}.jpg".format(video_filename[:-4], tid))
            cv2.imwrite(out_path, img)

    # release remaining snapshot track clips
    if generate_video_snapshots:
        for k, vid in vid_track_snapshots.items():
            vid.release()

    if save_static_trails:
        avg_frame = np.mean(avg_frame, axis=0)
        for tid, dots in static_track_trails.items():
            for dot in dots:
                cv2.circle(avg_frame, dot, 2, colors[tid], -1)
        #avg_frame = cv2.addWeighted(overlay, alpha, avg_frame, 1, 0)

        cv2.imwrite("%s_trails.jpg" % os.path.join(vis_dir, video_filename[:-4]), avg_frame)

def main(_argv):
    # Get a list of videos to visualise
    if os.path.exists(os.path.normpath(FLAGS.videos_dir)):
        videos = os.listdir(os.path.normpath(FLAGS.videos_dir))
        logging.info("Will process {} videos from {}".format(len(videos), os.path.normpath(FLAGS.videos_dir)))
    else:
        logging.info("videos_dir does not exist: {}".format(os.path.normpath(FLAGS.videos_dir)))
        return

    # generate frames
    for video in videos:
        visualise(os.path.join(os.path.normpath(FLAGS.videos_dir), video), FLAGS.frames_dir, FLAGS.detections_dir,
                  FLAGS.tracks_dir, FLAGS.stats_dir, FLAGS.vis_dir,
                  FLAGS.img_snapshots_dir, FLAGS.vid_snapshots_dir, FLAGS.around,
                  FLAGS.start_buffer, FLAGS.end_buffer,
                  FLAGS.display_tracks, FLAGS.display_detections, FLAGS.display_trails, FLAGS.save_static_trails,
                  FLAGS.generate_image_snapshots, FLAGS.generate_video_snapshots, FLAGS.summary, FLAGS.full)

if __name__ == '__main__':
    flags.DEFINE_string('videos_dir', 'data/unprocessed',
                        'Directory containing the video files to process')
    flags.DEFINE_string('frames_dir', 'data/frames',
                        'Directory to hold the frames as images')
    flags.DEFINE_string('detections_dir', 'data/detections',
                        'Directory to save the detection files')
    flags.DEFINE_string('tracks_dir', 'data/tracks',
                        'Directory to save the track files')
    flags.DEFINE_string('stats_dir', 'data/stats',
                        'Directory to hold the video stats')
    flags.DEFINE_string('vis_dir', 'data/summaries',
                        'Directory to hold the video visualisations')
    flags.DEFINE_string('img_snapshots_dir', 'data/snapshots/images',
                        'Directory to save image snapshots, if the flag --image_snapshots is used')
    flags.DEFINE_string('vid_snapshots_dir', 'data/snapshots/videos',
                        'Directory to save video snapshots, if the flag --video_snapshots is used')

    flags.DEFINE_string('around', 'detections',
                        'Base the shortening off of the detections or tracks?')
    flags.DEFINE_integer('start_buffer', 100,
                         'The number of frames to save pre-detection or track appearance.')
    flags.DEFINE_integer('end_buffer', 50,
                         'The number of frames to save post-detection or track appearance.')

    flags.DEFINE_boolean('display_tracks', False,
                         'Do you want to save a video with the tracks?')
    flags.DEFINE_boolean('display_detections', False,
                         'Do you want to save a video with the detections?')
    flags.DEFINE_boolean('display_trails', False,
                         'Do you want display trails after the tracks?')
    flags.DEFINE_boolean('save_static_trails', False,
                         'Do you want to save an mean image with all track trails printed?')
    flags.DEFINE_boolean('generate_image_snapshots', True,
                         'Do you want to save image snapshots for each track?')
    flags.DEFINE_boolean('generate_video_snapshots', False,
                         'Do you want to save video snapshots for each track?')
    flags.DEFINE_boolean('summary', False,
                         'Do you want to save out a summary video?')
    flags.DEFINE_boolean('full', False,
                         'Do you want to save out the full video?')

    try:
        app.run(main)
    except SystemExit:
        pass


"""add unique constraint

Revision ID: 49bb5a76c589
Revises: 17615e42022c
Create Date: 2018-03-04 09:49:13.381590

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '49bb5a76c589'
down_revision = '17615e42022c'
branch_labels = None
depends_on = None

def upgrade():
    conn = op.get_bind()
    conn.execute(
        sa.sql.text(
            """
            ALTER TABLE stocks ADD CONSTRAINT stocks_label_uq UNIQUE (label);
            """
        )
    )

def downgrade():
    conn = op.get_bind()
    conn.execute(
        sa.sql.text(
            """
            ALTER TABLE stocks DROP CONSTRAINT stocks_label_uq;
            """
        )
    )