import json
import uuid

from scrapy.http import HtmlResponse

from qidian.items import *

class WanbenSpider(scrapy.Spider):
    name = 'wanben'
    allowed_domains = []
    start_urls = ['https://www.qidian.com/finish']
    token = ''

    def parse(self, response: HtmlResponse):
        self.token = dict(response.headers).get(b'Set-Cookie')[0].decode('utf-8').split(';')[0].split('=')[-1]

        print(self.token)
        print('sdddd')
        if response.status == 200:
            lis = response.css('.all-img-list li')
            for li in lis:
                item = BookItem()
                item['book_id'] = uuid.uuid4().hex
                a = li.xpath('.//div[@class="book-img-box"]/a')
                item['book_url'] = "https:{}".format(a.xpath('./@href').get())
                item['book_cover'] = "https:{}".format(a.xpath('./img/@src').get())
                item['book_name'] = li.xpath('.//div[2]/h4//text()').get()
                item['author'], *item['tags'] = li.css('.author a::text').extract()
                item['summary'] = li.css('.intro::text').get().strip()
                book_singal = item['book_url'].split('/')[-1]
                detail_url = "https://book.qidian.com/ajax/book/category?_csrfToken={}&bookId={}".format(self.token,
                                                                                book_singal)

                print(detail_url)
                # 请求详情页章节
                yield scrapy.Request(detail_url, callback=self.parse_info,
                                     meta={'book_id': item['book_id']})

                yield item
            next_page = response.css('.lbf-pagination-item-list').xpath('./li[last()]/a/@href').get()
            if next_page.find('javascript') == -1:  # 存在下一页
                yield scrapy.Request('https:{}'.format(next_page))

    def parse_info(self, response: HtmlResponse):
        """解析小说的详情页章节"""
        book_id = response.meta['book_id']
        chapters = json.loads(response.text).get("data").get('vs')

        for chapter in chapters:
            cs = chapter.get('cs')
            for data in cs:
                item = SegItem()
                item['seg_id'] = uuid.uuid4().hex
                item['book_id'] = book_id
                item['url'] = "https://read.qidian.com/chapter/{}".format(data.get('cU'))  # 章节详情标识
                item['title'] = data.get('cN')  # 章节标题
                yield scrapy.Request(item['url'], callback=self.parse_seg, meta={'seg_id': item['seg_id']})
                yield item

    def parse_seg(self, response):
        """解析章节的小说内容"""
        print("-" * 50, "进入章节内容", "-" * 50)
        item = SegDetailItem()
        item['seg_id'] = response.meta['seg_id']
        contents = "".join(response.css('.read-content p::text').extract())
        item['content'] = contents.replace('\u3000', '').replace('\n', '').strip()
        yield item


# Classe principal do código
#
# exec -> python runTests.py caminha_ate_a_pasta_do_projeto
#
# <b>runTests</b>: Esse script controla a execução de todas as partes do algoritmo.
# É nele que são lidas as imagens e estas são passadas para as demais etapadas do algoritmo
# Ao término no processo, ele gera arquivos .csv (tabelas) com os resultados individuais de cada imagem
# e a média geral de todas as execuções.
#
# Explicando resumidamente o código:
#
#
# author:
# Pablo Vinicius - https://github.com/pabloVinicius
#
# contributors:
# Filipe A. Sampaio - https://github.com/filipeas

# imports dos arquivos necessários
from main import vsf # responsável pela separação dos superpixels
from features_extraction import ftet # responsável por extrair as caracteristicas de cada superpixel usadas no modelo
from select_and_classify import classify, select_random_seeds # responsável pela classificação do modelo criado
from statistics import mean, pvariance
import os, csv, time
import numpy as np
import sys

def index(original, marcada):
    # Porcentagens de dados das imagens usadas para treino.
    percentages = [
        # 0.01,
        # 0.05,
        # 0.1,
        # 0.15,
        # 0.2,
        0.25,
        # 0.3,
        # 0.35,
        # 0.4,
        # 0.5
    ]

    # Quantidade de segmentos para o superpixel
    qtdSegments = [
        1500,
        # 2000,
        # 2500,
        # 3000,
        # 4000,
        # 5000
    ]

    # Tenta encontrar o diretório results. Se ele não existir, o cria.
    results_path = 'results/'
    if not os.path.exists(results_path):
        os.makedirs(results_path)

    # Cria a estrutura de dados que irá guardar os resultados das médias das execuções para todas as imagens e para todas as métricas.
    # Funciona como uma matriz  de (Nª de porcentagenx utilizadas)x(número de métricas, que no caso são 4).
    # No final, essa estrutura é utilizada para calcular a média das médias de cada porcentagem de dados pra treino.
    total_means = [[[],[],[],[]] for i in range(len(percentages))]

    # instanciando variavel pra guardar todas as acuracias da imagem corrente
    acuracia_corrent = [[[],[],[],[],[]] for i in range(len(percentages))]

    # instanciando variavel pra guardar todas as sensibilidade da imagem corrente
    sensibilidade_corrent = [[[],[],[],[],[]] for i in range(len(percentages))]

    # instanciando variavel pra guardar todas as especificidade da imagem corrente
    especificidade_corrent = [[[],[],[],[],[]] for i in range(len(percentages))]

    # instanciando variavel pra guardar todas as dice da imagem corrente
    dice_corrent = [[[],[],[],[],[]] for i in range(len(percentages))]

    print(f'\n\n########### Executando a imagem ############\n\n')

    # Marcando o tempo de execução.
    start = time.time()

    # Abrindo a imagem original e a imagem marcada.
    path = f'{original}' # original
    path_marked = f'{marcada}' # marcada

    print(path + '\n' + path_marked)

    # iterando sobre a quantidade de segmentos de imagem para o superpixel
    for segment in qtdSegments:
        # Executa a primeira etapa do algoritmo, responsável pela separação dos superpixels da imagem (arquivo main.py).
        vsf(path, path_marked, segment)

        # Executa a segunda etapa do algoritmo, responsável pela extração de características de cada superpixel gerado na etapa anterior.
        # Disponível no arquivo features_extraction.py
        ftet()

        # Para cada percentual:
        for index, percent in enumerate(percentages):

            # Cria uma estrutura de dados para salvar os resultados e posteriormente calcular as médias das métricas
            metrics_media = [[], [], [], []]

            # Executa 5 vezes para depois tirar a média (5 vezes pois são 5 métricas diferentes)
            for i in range(1):

                # Terceira etapa do algoritmo:
                # Usa as características extraídas na etapa anterior para classificar e gerar as métricas para cada imagem
                # Disponível no arquivo select_and_classify.py
                acc, sen, spe, dice = classify(percent) # classificação por Random Forest
                # acc, sen, spe, dice = select_random_seeds(image, percent) # classificação por sfc-means

                # Coloca as métricas na estrutura de dados
                metrics_media[0].append(acc)
                metrics_media[1].append(sen)
                metrics_media[2].append(spe)
                metrics_media[3].append(dice)

            # Depois adiciona na estrutura de dados os dados gerais
            total_means[index][0].append(mean(metrics_media[0])*100) # acuracia
            total_means[index][1].append(mean(metrics_media[1])*100) # sensibilidade
            total_means[index][2].append(mean(metrics_media[2])*100) # especificidade
            total_means[index][3].append(mean(metrics_media[3])) # dice

            if(segment == 1500):
                acuracia_corrent[index][0].append(mean(metrics_media[0])*100)
                sensibilidade_corrent[index][0].append(mean(metrics_media[1])*100)
                especificidade_corrent[index][0].append(mean(metrics_media[2])*100)
                dice_corrent[index][0].append(mean(metrics_media[3])*100)
            elif(segment == 2000):
                acuracia_corrent[index][1].append(mean(metrics_media[0])*100)
                sensibilidade_corrent[index][1].append(mean(metrics_media[1])*100)
                especificidade_corrent[index][1].append(mean(metrics_media[2])*100)
                dice_corrent[index][1].append(mean(metrics_media[3])*100)
            elif(segment == 2500):
                acuracia_corrent[index][2].append(mean(metrics_media[0])*100)
                sensibilidade_corrent[index][2].append(mean(metrics_media[1])*100)
                especificidade_corrent[index][2].append(mean(metrics_media[2])*100)
                dice_corrent[index][2].append(mean(metrics_media[3])*100)
            elif(segment == 3000):
                acuracia_corrent[index][3].append(mean(metrics_media[0])*100)
                sensibilidade_corrent[index][3].append(mean(metrics_media[1])*100)
                especificidade_corrent[index][3].append(mean(metrics_media[2])*100)
                dice_corrent[index][3].append(mean(metrics_media[3])*100)
            elif(segment == 4000):
                acuracia_corrent[index][4].append(mean(metrics_media[0])*100)
                sensibilidade_corrent[index][4].append(mean(metrics_media[1])*100)
                especificidade_corrent[index][4].append(mean(metrics_media[2])*100)
                dice_corrent[index][4].append(mean(metrics_media[3])*100)

    # calcula o tempo de execução do algoritmo para a imagem atual.
    end = time.time()
    print(f'\n\n########### Fim da execucao da imagem no tempo de {end-start} segundos ############\n\n')
    print("Result = ", acuracia_corrent[0][0], sensibilidade_corrent[0][0], especificidade_corrent[0][0], dice_corrent[0][0], " Fim")
    return acuracia_corrent[0][0], sensibilidade_corrent[0][0], especificidade_corrent[0][0], dice_corrent[0][0]


from sqlalchemy import MetaData, Table, String, Integer, Column, Boolean, DateTime
from sqlalchemy import create_engine, insert, update, delete, select, exc

from ..Telnet.Alarm import Alarm

class AlarmDatabase:
    def __init__(self):
        self.engine = create_engine('sqlite:///alarms.db')
        self.create_alarm_table()
        self.clear_tables()

    def clear_tables(self):
        conn = self.engine.connect()
        conn.execute(delete(self.alarms))
        conn.execute(delete(self.nodes))

    def create_alarm_table(self):
        metadata = MetaData()
        self.alarms = Table('alarmwindow_alarms', metadata,
                            Column('id', Integer(), nullable=False),
                            Column('type', String(7), nullable=False),
                            Column('raising_time', DateTime, nullable=False),
                            Column('ceasing_time', DateTime, nullable=True),
                            Column('managed_object', String(30)),
                            Column('object_name', String(30)),
                            Column('slogan', String(30)),
                            Column('descr', String(60)),
                            Column('text', String(300), nullable=False),
                            Column('is_active', Boolean(), nullable=False),
                            Column('node_id', Integer()),
                            Column('node_update_id', Integer())
                            )

        self.nodes = Table('alarmwindow_nodes', metadata,
                           Column('id', Integer(), autoincrement=True, primary_key=True),
                           Column('name', String(20), nullable=False, unique=True),
                           Column('update_id', Integer(), default=0)
                           )
        metadata.create_all(self.engine)

    def get_current_update_id(self, node_id) -> int:
        conn = self.engine.connect()
        query = select(self.nodes.c.update_id).where(
            self.nodes.c.id == node_id
        )
        result = conn.execute(query).fetchone()
        return result[0]

    def insert_new_alarms(self, alarm_objects: list[Alarm]):
        if not len(alarm_objects):
            return
        target_node = None
        update_id = self.get_current_update_id(alarm_objects[0].node_id)
        query = insert(self.alarms).values(
            [
                {
                    'id': alarm.id,
                    'type': alarm.type,
                    'raising_time': alarm.raising_time,
                    'managed_object': alarm.managed_object,
                    'object_name': alarm.object_name,
                    'slogan': alarm.slogan,
                    'descr': alarm.descr,
                    'text': alarm.text,
                    'is_active': alarm.is_active,
                    'node_id': alarm.node_id,
                    'node_update_id': update_id if alarm.node_id == target_node
                    else self.get_current_update_id(alarm.node_id)
                } for alarm in alarm_objects
            ]
        )

        conn = self.engine.connect()
        conn.execute(query)

    def update_ceased_alarms(self, alarm_objects: list[Alarm]):
        if not len(alarm_objects):
            return
        conn = self.engine.connect()
        transaction = conn.begin()
        last_node = None
        update_id = self.get_current_update_id(alarm_objects[0].node_id)
        for alarm in alarm_objects:
            if last_node != alarm.node_id:
                update_id = self.get_current_update_id(alarm.node_id)
                last_node = alarm.node_id
            upd = update(self.alarms).where(
                self.alarms.c.id == alarm.id
            ).values({
                'is_active': False,
                'ceasing_time': alarm.ceasing_time,
                'node_update_id': update_id
            })
            conn.execute(upd)
        transaction.commit()

    def increase_update_id(self, controller_id):
        conn = self.engine.connect()
        upd = update(self.nodes).where(
            self.nodes.c.id == controller_id
        ).values(
            {'update_id': self.nodes.c.update_id + 1}
        )
        conn.execute(upd)

    def add_node(self, node_name):
        conn = self.engine.connect()
        query = insert(self.nodes).values(
            {
                'name': node_name,
                'update_id': 0
            }
        )
        try:
            conn.execute(query)
        except exc.IntegrityError:
            pass

    def get_node_id(self, name) -> int:
        conn = self.engine.connect()
        query = select(self.nodes.c.id).where(
            self.nodes.c.name == name
        )
        result = conn.execute(query).fetchone()
        return result[0]


import pytest
import string
import getpass
from playwright.sync_api import Page, Playwright
from test_commonmethod import test_open_page_sso,test_create_ppe_codespace,test_createAndinstall
from test_delete_codespace import test_deleteAllCodespace

@pytest.mark.daily
@pytest.mark.disconnect
def test_delete_all_codespaces(playwright: Playwright):
    test_deleteAllCodespace(playwright)

@pytest.mark.daily
@pytest.mark.disconnect
def test_disconnect_ViewCreationLog_codespace(playwright: Playwright):
    pageurl="https://github.com/codespaces/new?location=EastUs"
    page=test_open_page_sso(playwright,pageurl)
    try:
        page.context.pages[-2].close()
        test_create_ppe_codespace(page, "Microsoft/vscode-remote-try-python")
        test_createAndinstall(page, "16-core")
        #View Creation Log command
        page.wait_for_timeout(15000)
        page.keyboard.press("Control+Shift+P")
        page.keyboard.type("View Creation Log")
        page.keyboard.press("ArrowDown")
        page.keyboard.press("Enter")
        page.wait_for_timeout(1500)
        if not page.get_by_role("tab").get_by_title("/workspaces/.codespaces/.persistedshare/creation.log").is_visible():
            page.wait_for_timeout(4000)
            page.keyboard.press("Control+Shift+P")
            page.keyboard.type("View Creation Log")
            page.keyboard.press("ArrowDown")
            page.keyboard.press("Enter")
            page.wait_for_timeout(1500)
        assert page.get_by_role("tab").get_by_title("/workspaces/.codespaces/.persistedshare/creation.log").is_visible()
        #disconnect codespace
        page.keyboard.press("Control+Shift+P")
        page.keyboard.type("Codespaces: details")
        page.keyboard.press("ArrowDown")
        page.keyboard.press("Enter")
        page.wait_for_timeout(1500)
        page.get_by_role("tree").nth(0).hover()
        page.get_by_title("Disconnect").click()
        page.wait_for_timeout(1500)
        assert 'https://github.com/codespaces/' in page.url
    finally:
        page.close()


class SetCover():
    # initializes the object with the dataset, parses the data
    def __init__(self, inputFile, printData=False):
        with open(inputFile) as f:
            content = f.read()
            lines = content.split("\n")
            firstLine = lines[0].split()
            # m is the number of rows, which correspond to elements
            self.m = int(firstLine[0])
            # n is the number of columns, which correspond to sets
            self.n = int(firstLine[1])
            # sets is a list where sets[i] corresponds to the weight of the i-th set
            sets = []
            # elements is a list of lists where elements[i] is the set of sets that contain element i
            elements = []
            temp = []
            for i in range(1, len(lines)):
                line = lines[i].split()
                if len(sets) < self.n:
                    for weight in line:
                        sets.append(int(weight))
                elif len(line) == 1:
                    if len(temp) != 0:
                        elements.append(temp)
                        temp = []
                    else:
                        continue
                else:
                    for setNumber in line:
                        temp.append(int(setNumber))
            elements.append(temp)
            self.sets = sets
            self.elements = elements
            if printData:
                print(f"m is {self.m}")
                print(f"n is {self.n}")
                print(self.sets)
                print(self.elements)

    # the Integer Linear Program version of Set Cover, solved exactly using the CP-SAT ILP Solver
    def solveILP(self, timeLimit=None, printData=False):
        from ortools.sat.python import cp_model
        model = cp_model.CpModel()
        elements = self.elements
        sets = self.sets

        # indicator variables for whether or not the i-th set is included in the solution
        set_vars = [model.NewBoolVar(f'indicator for set {i}') for i in range(self.n)]

        # constraint that each vertex needs to be covered
        for element in elements:
            setCoverConstraint = 0
            for elementSet in element:
                setCoverConstraint += set_vars[elementSet - 1]
            model.Add(setCoverConstraint >= 1)

        # minimization of the total weight
        total_weight = 0
        for i in range(len(sets)):
            total_weight += (sets[i] * set_vars[i])
        model.Minimize(total_weight)

        # solve the constraint program and print the results
        solver = cp_model.CpSolver()
        if timeLimit is not None:
            solver.parameters.max_time_in_seconds = timeLimit
        from timeit import default_timer as timer
        t = timer()
        result = solver.Solve(model)
        if result == cp_model.OPTIMAL:
            answer = 0
            for j in range(len(sets)):
                answer += (solver.Value(set_vars[j]) * sets[j])
            if printData:
                print(f'the ILP scheme takes time {timer() - t} seconds')
                print(f'the optimal set cover weight is {answer}')
            return (timer() - t, answer)
        elif result == cp_model.FEASIBLE:
            answer = 0
            for j in range(len(sets)):
                answer += (solver.Value(set_vars[j]) * sets[j])
            if printData:
                print(f'the ILP scheme takes time {timer() - t} seconds')
                print(f'the feasible set cover weight is {answer}')
            return (timer() - t, answer)
        else:
            print('this instance of Set Cover is not solvable')

    # the Linear Programming + Randomized Rounding Scheme version of Set Cover, approximated by using the GLOP solver and rounding scheme
    def solveLP_Rounding(self, printData=False):
        from ortools.linear_solver import pywraplp
        import random
        solver = pywraplp.Solver('set_cover_LP_solver', pywraplp.Solver.GLOP_LINEAR_PROGRAMMING)
        elements = self.elements
        sets = self.sets
        # indicator variables for whether or not the i-th set is included in the solution
        set_vars = [solver.NumVar(0, 1, f'indicator for set {i}') for i in range(self.n)]
        # constraint that each vertex needs to be covered
        for element in elements:
            setCoverConstraint = 0
            for elementSet in element:
                setCoverConstraint += set_vars[elementSet - 1]
            solver.Add(setCoverConstraint >= 1)
        # minimization of the total weight using linear programming
        total_weight = 0
        for i in range(len(sets)):
            total_weight += (sets[i] * set_vars[i])
        solver.Minimize(total_weight)

        # solve the constraint program
        from timeit import default_timer as timer
        t = timer()
        if solver.Solve() == pywraplp.Solver.OPTIMAL:
            set_answers = [set_vars[i].solution_value() for i in range(len(set_vars))]
        else:
            print("impossible to solve this instance of set cover!")
            return

        # randomzied rounding scheme: add the set to the solution with probability corresponding to the value of the variable. Run this log n times
        answer = 0
        counter = 1
        added = set()
        while (counter < self.n):
            for setElem in range(len(sets)):
                if random.random() < set_answers[setElem]:
                    added.add(setElem)
            counter = counter * 2
        for addedSet in added:
            answer += sets[addedSet]

        # for elements that are not covered after the rounding scheme, add the minimum weight set that covers it
        for element in elements:
            covered = False
            for elementSet in element:
                if (elementSet - 1) in added:
                    covered = True
                    break
            if not covered:
                elementWeights = [sets[i - 1] for i in element]
                answer += min(elementWeights)
        # print the results
        if printData:
            print(f'the LP Rounding scheme takes time {timer() - t} seconds')
            print(f'the optimal set cover weight is {answer}')

        return (timer() - t, answer)


#!/usr/bin/env python3
# _coding_: utf8

import json
import sys
import time
import traceback
import random as rd
import pandas as pd
from threading import Thread
from broker.client import BrokerClient

class icoMarketMaking(object):

    def __init__(self, symbol, pm):
        self.symbol = symbol
        self.pm = pm

        self.api = BrokerClient(entry_point=self.pm.get('entry_point', 'https://api.bhex.com/openapi/'), api_key=self.pm['api_key'], secret=self.pm['secret'])

        # self.api = BhexClient(api_key=self.pm['api_key'], secret=self.pm['secret'], entry_point=self.pm.get('entry_point', 'https://api.bhex.com/openapi/'))
        self.imitate_symbol = self.pm.get('imitate_symbol', '')

        # load order config file
        self.tgtStructBid = pd.DataFrame.from_dict(self.pm['tgtStructBid'])
        self.tgtStructAsk = pd.DataFrame.from_dict(self.pm['tgtStructAsk'])
        # print(self.tgtStructBid.loc[str(1), 'fromBestBid'], self.tgtStructAsk)

        self.orderThread = None
        self.klineThread = None

    def start(self):
        orders = self.api.open_orders(symbol=self.symbol, limit=1000)
        self.tgtStructBid['cid'] = [''] * len(self.tgtStructBid)
        self.tgtStructAsk['cid'] = [''] * len(self.tgtStructAsk)

        orders = sorted(orders, key=lambda x: float(x['price']), reverse=True)
        i = 0
        for o in orders:
            if o['side'] == 'BUY' and o['clientOrderId'][:len(self.pm['orderPrefix'])] == self.pm['orderPrefix']:
                if i < self.pm['orderNum']:
                    self.tgtStructBid.loc[str(i), 'cid'] = o['clientOrderId']
                    i += 1
                else:
                    self.api.order_cancel(client_order_id=o['clientOrderId'])

        i = 0
        for o in reversed(orders):
            if o['side'] == 'SELL' and o['clientOrderId'][:len(self.pm['orderPrefix'])] == self.pm['orderPrefix']:
                if i < self.pm['orderNum']:
                    self.tgtStructAsk.loc[str(i), 'cid'] = o['clientOrderId']
                    i += 1
                else:
                    self.api.order_cancel(client_order_id=o['clientOrderId'])

        self.orderThread = Thread(target=self.orderWorker)
        self.orderThread.start()

        self.klineThread = Thread(target=self.drawKline)
        self.klineThread.start()

    # to calculate the fair bid/ask price
    def getFairPrice(self):
        try:
            depth = self.api.depth(symbol=self.symbol, limit=100)
            print(str(depth['bids']))
            bids = depth['bids']
            pqsum = qsum = 0
            for b in bids:
                qsum += float(b[1])
                pqsum += float(b[0]) * float(b[1])
                if qsum >= self.pm['fairBidPriceQty']:
                    break
            fair_bid_price = pqsum / qsum

            asks = depth['asks']
            pqsum = qsum = 0
            for a in asks:
                qsum += float(a[1])
                pqsum += float(a[0]) * float(a[1])
                if qsum >= self.pm['fairAskPriceQty']:
                    break
            fair_ask_price = pqsum / qsum

            if self.imitate_symbol:
                tgt = self.api.klines(symbol=self.imitate_symbol, limit=1)[0]
                chg = float(tgt[4]) / float(tgt[1])
                fair_bid_price *= chg
                fair_ask_price *= chg

            return fair_bid_price, fair_ask_price
        except:
            traceback.print_exc()

    # order placing fn, will check one pair of bid and ask at one time
    def orderWorker(self):
        while True:
            for i in reversed(range(self.pm['orderNum'])):
                print(i)
                try:
                    fair_bid_price, fair_ask_price = self.getFairPrice()
                    # print(fair_bid_price, fair_ask_price)
                    fair_mid_price = 0.5 * (fair_bid_price + fair_ask_price)
                    best_bid = fair_mid_price * (1 - self.pm['bid_spread'])
                    best_ask = fair_mid_price * (1 + self.pm['ask_spread'])

                    send_flag = True
                    price = best_bid * (1 - self.tgtStructBid.loc[str(i), 'fromBestBid'])
                    order = self.api.order_get(orig_client_order_id=self.tgtStructBid.loc[str(i), 'cid'])
                    print(str(order))
                    try:
                        if order['status'] == 'FILLED':
                            None
                        elif abs(float(order['price']) - price) > 0.0005 * price:
                            self.api.order_cancel(client_order_id=self.tgtStructBid.loc[str(i), 'cid'])
                        else:
                            send_flag = False
                    except:
                        traceback.print_exc()
                        print(order)

                    if send_flag:
                        params = {'symbol': self.symbol, 'side': 'BUY', 'type': 'LIMIT_MAKER',
                                  'quantity': self.pm['orderQty'] * self.tgtStructBid.loc[str(i), 'qtyMultiplier'] * (
                                          1 + 0.6 * (rd.random() - 0.5)),
                                  'price': price * (1 + 0.0005 * (rd.random() - 0.5)),
                                  'newClientOrderId': self.pm['orderPrefix'] + str(int(time.time() * 1000)) + 'b' + str(i)}
                        self.prepareOrder(params)
                        self.tgtStructBid.loc[str(i), 'cid'] = params['newClientOrderId']

                    send_flag = True
                    price = best_ask * (1 + self.tgtStructAsk.loc[str(i), 'fromBestAsk'])
                    order = self.api.order_get(orig_client_order_id=self.tgtStructAsk.loc[str(i), 'cid'])
                    try:
                        if order['status'] == 'FILLED':
                            None
                        elif abs(float(order['price']) - price) > 0.0005 * price:
                            self.api.order_cancel(client_order_id=self.tgtStructAsk.loc[str(i), 'cid'])
                        else:
                            send_flag = False
                    except:
                        traceback.print_exc()
                        print(order)

                    if send_flag:
                        params = {'symbol': self.symbol, 'side': 'SELL', 'type': 'LIMIT_MAKER',
                                  'quantity': self.pm['orderQty'] * self.tgtStructAsk.loc[str(i), 'qtyMultiplier'] * (
                                          1 + 0.6 * (rd.random() - 0.5)),
                                  'price': price * (1 + 0.0005 * (rd.random() - 0.5)),
                                  'newClientOrderId': self.pm['orderPrefix'] + str(int(time.time() * 1000)) + 'a' + str(i)}
                        self.prepareOrder(params)
                        self.tgtStructAsk.loc[str(i), 'cid'] = params['newClientOrderId']

                except:
                    traceback.print_exc()

                time.sleep(self.pm.get('orderRefreshInterval', 5))

    # draw k line fn
    def drawKline(self):
        while True:
            try:
                res = self.api.ticker_24hr(symbol=self.symbol)
                # print(res)
                best_bid = float(res['bestBidPrice'])
                best_ask = float(res['bestAskPrice'])

                cid = str(int(time.time()))

                order = {'symbol': self.symbol,
                         'quantity': self.pm['orderQty'] * (0.2 + 1.6 * rd.random()),
                         'price': round(best_bid + (best_ask - best_bid) * rd.random(), self.pm['priceDecimalLength']),
                         'newClientOrderId': cid}

                qty = order['quantity']

                if best_bid + 0.5 * (best_ask - best_bid) < order['price'] < best_ask or order['price'] == best_ask:
                    order['side'] = 'SELL'
                    order['type'] = 'LIMIT_MAKER'
                    self.prepareOrder(order)
                    # time.sleep(3 + 7 * rd.random())
                    order.pop('newClientOrderId')
                    order['side'] = 'BUY'
                    order['type'] = 'LIMIT'
                    order['timeInForce'] = 'IOC'
                    order['quantity'] = qty * (0.1 + 0.8 * rd.random())
                    self.prepareOrder(order)
                    # time.sleep(2 + 5 * rd.random())
                    if rd.random() < 0.2:
                        order['quantity'] = qty - order['quantity']
                        self.prepareOrder(order)
                    self.api.order_cancel(client_order_id=cid)

                elif best_bid < order['price'] < best_bid + 0.5 * (best_ask - best_bid) or order['price'] == best_bid:
                    order['side'] = 'BUY'
                    order['type'] = 'LIMIT_MAKER'
                    self.prepareOrder(order)
                    # time.sleep(3 + 7 * rd.random())
                    order.pop('newClientOrderId')
                    order['side'] = 'SELL'
                    order['type'] = 'LIMIT'
                    order['timeInForce'] = 'IOC'
                    order['quantity'] = qty * (0.1 + 0.8 * rd.random())
                    self.prepareOrder(order)
                    # time.sleep(2 + 5 * rd.random())
                    if rd.random() < 0.2:
                        order['quantity'] = qty - order['quantity']
                        self.prepareOrder(order)
                    self.api.order_cancel(client_order_id=cid)

            except:
                traceback.print_exc()
                print(res)

            time.sleep(10 + 30 * rd.random())

    # random cancel order fn
    def randomCancel(self):
        try:
            orders = self.api.open_orders(symbol=self.symbol, limit=1000)
            i = int(rd.random() * len(orders))
            self.api.order_cancel(client_order_id=orders[i]['clientOrderId'])
            print('randomCancel', orders[i])
        except:
            traceback.print_exc()

    # send order fn
    def prepareOrder(self, params):
        params['price'] = round(params['price'], self.pm['priceDecimalLength'])
        params['quantity'] = round(params['quantity'], self.pm['qtyDecimalLength'])
        res = self.api.order_new(**params)
        if res.get('msg', '') == 'Balance insufficient ':
            print(res, params)
            self.randomCancel()
            print(self.api.account())
        elif 'msg' in res:
            print(res, params)
        # else:
        #     print(res, params)

if __name__ == '__main__':

    # try:
    #     symbol = str.upper(sys.argv[1])
    # except:
    #     symbol = ''
    #     print("NEED SYMBOL")
    #     exit(0)

    symbol = 'EEGUSDT'

    with open("parameter.json") as f:
        json_file = json.load(f)

    a = icoMarketMaking(symbol=symbol,
                        pm=json_file[symbol])

    a.start()


import os
import cv2
from datetime import date, datetime
import pandas as pd
from tensorflow.image import ssim
import numpy as np
import tensorflow as tf
from google.cloud import storage
from tempfile import TemporaryFile, NamedTemporaryFile
import matplotlib.pyplot as plt

BUCKET_NAME = 'm_digital_dream_bucket'
client = storage.Client()
bucket = client.bucket(BUCKET_NAME)

today = date.today().isoformat()
time = datetime.now().strftime("%H-%M")
IMAGE = '%d.jpg'

dir_1 = ('raw_data')
path = f'gs://{BUCKET_NAME}/Clusters/dossier_csv'
out_put_path = ('Clusters')
clusters = ['cluster_3_5_14_.csv',
            'cluster_pca_tsne_3_5_19.csv',
            'cluster_2021-06-15_14-08.csv',
            'cluster_2021-06-15_14-11.csv',
            'cluster_2021-06-15_14-13.csv',
            'cluster_2021-06-15_14-41.csv',
            'cluster_2021-06-15_14-48.csv',
            'cluster_2021-06-15_14-49.csv']

# def pour comparé les images cote à cote et retourne le SSIM
def compare_images(imageA, imageB):
    im1 = tf.expand_dims(imageA, axis=0)
    im2 = tf.expand_dims(imageB, axis=0)
    im1_resize = tf.reshape(im1, (1, 512, 512, 3))
    im2_resize = tf.reshape(im2, (1, 512, 512, 3))
    s = ssim(im1_resize, im2_resize, max_val=255)
    return s

# def pour lire les images d'origines en couleur
def read_gcp(bucket, path):
    blob = bucket.blob(path)
    image = np.asarray(bytearray(blob.download_as_string()), dtype="uint8")
    imageBGR = cv2.imdecode(image, cv2.IMREAD_COLOR)
    return imageBGR

# boucle pour comparé les noms d'images de clusters et les images d'origines
# les csv de nom d'image des clusters
for i, cluster in enumerate(clusters):
    df = pd.read_csv(f'{path}/{cluster}')
    #print(i, cluster)
    #print(len(df['filename'])-1)
    incre = 0
    # regarde image par image pour retrouvé les images d'origine
    for j in range(len(df['filename'])-1):
        if incre < 19:
            #print(os.path.join(dir_1, df['filename'][j]))
            pic1 = read_gcp(bucket ,os.path.join(dir_1, df['filename'][j])) # range la bonne image dans la var pic1 et rajoute le RGB
            pic1 = cv2.resize(pic1,(512,512)) # adapte l'image en 512, 512
            pic2 = read_gcp(bucket, os.path.join(dir_1, df['filename'][j+1]))
            pic2 = cv2.resize(pic2,(512,512))
            s = compare_images(pic1, pic2) # retourne le SSIM entre les 2 images
            #print(j)
            #print(s)
            if s < 0.9: # si SSIM < 0.9 alors les 2 images sont pas identique donc rajout des 2 images
                name = IMAGE % incre
                with NamedTemporaryFile(suffix='.jpg') as gcs_image:
                    # recupére le nom de l'image et le lien dans gcs_image
                    cv2.imwrite(gcs_image.name, pic1)
                    gcs_image.seek(0)
                    # créer un dossier pour les images du cluster i
                    blob = bucket.blob(f'{out_put_path}/cluster_{today}_{time}_{i}/{name}')
                    # upload le gcs_image
                    blob.upload_from_file(gcs_image)
                incre = incre+1
                name = IMAGE % incre
                with NamedTemporaryFile(suffix='.jpg') as gcs_image:
                    cv2.imwrite(gcs_image.name, pic2)
                    gcs_image.seek(0)
                    blob = bucket.blob(f'{out_put_path}/cluster_{today}_{time}_{i}/{name}')
                    blob.upload_from_file(gcs_image)
            else: # si SSIM > 0.9 alors les 2 images sont identique on ne récupére que la 2éme image
                name = IMAGE % incre
                with NamedTemporaryFile(suffix='.jpg') as gcs_image:
                    cv2.imwrite(gcs_image.name, pic2)
                    gcs_image.seek(0)
                    blob = bucket.blob(f'{out_put_path}/cluster_{today}_{time}_{i}/{name}')
                    blob.upload_from_file(gcs_image)
                incre = incre-1


import requests

from setup import config

app_token = config['PUSHOVER_APP_TOKEN']
user_key = config['PUSHOVER_USER_KEY']

api_endpoint = 'https://api.pushover.net/1/messages.json'

def send_notification(title, message) -> requests.Response:
    parameters = {
        'token': app_token,
        'user': user_key,
        'message': message,
        'title': title
    }

    response = requests.post(api_endpoint, parameters)
    return response


from django.contrib import admin
from .models import Trainer

class TrainerAdmin(admin.ModelAdmin):
    list_display = (
        'name',
        'years_practicing',
        'years_teaching',
        'bio',
        'image'
    )

    ordering = ('name'),

# Register your models here.
admin.site.register(Trainer, TrainerAdmin)


#! usr/bin/env/python
'''
	Data download from Yahoo Finance

	Date,Open,High,Low,Close,Volume,Adj Close
'''

import numpy as np
import sys
from matplotlib.pyplot import plot
from matplotlib.pyplot import show

N = 8	# window size
weights = np.hanning(N)
print("Weights {0}".format(weights))

avgo = np.loadtxt('AVGO_clipped.csv', delimiter=',', usecols=(4,), unpack=True)
avgo = avgo[::-1]

qrvo = np.loadtxt('QRVO_clipped.csv', delimiter=',', usecols=(4,), unpack=True)
qrvo = qrvo[::-1]

# return
avgo_returns = np.diff(avgo)/avgo[:-1]
qrvo_returns = np.diff(qrvo)/qrvo[:-1]

# smooth
smooth_avgo = np.convolve(weights/weights.sum(), avgo_returns)[N-1:-N+1]
smooth_qrvo = np.convolve(weights/weights.sum(), qrvo_returns)[N-1:-N+1]

t = np.arange(N-1, len(avgo_returns))

plot(t, avgo_returns[N-1:], lw = 1.0)
plot(t, smooth_avgo, lw=2.0)
plot(t, qrvo_returns[N-1:], lw = 1.0)
plot(t, smooth_qrvo, lw = 2.0)
show()

K = 3
poly_avgo = np.polyfit(t, smooth_avgo, K)
poly_qrvo = np.polyfit(t, smooth_qrvo, K)

poly_sub = np.polysub(poly_avgo, poly_qrvo)

xpoints = np.roots(poly_sub)
print("Intersection points {0}".format(xpoints))

reals = np.isreal(xpoints)
print("Real number? {0}".format(reals))

xpoints = np.select([reals], [xpoints])
xpoints = xpoints.real
print("Real intersection points {0}".format(xpoints))

print("Sans 0s {0}".format(np.trim_zeros(xpoints)))


from PIL import Image, ImageColor, ImageDraw
from cmath import *
from math import *
# Note: output images will be saved to the same directory where you
# store this .py file.
# images will not pop up, will just be saved.
# initialize an image
xmin = -5/2
xmax = 5/2
ymin = -5/2
ymax = 5/2
w = 500
h = 500
bitmap = Image.new("RGB", (w, h), "white")
pix = bitmap.load()
deltax = abs(xmax - xmin)/w
deltay = abs(ymax - ymin)/h

# input z and c, out put the greens function value
def green_function(zx0, zy0, c):
    if zx0 == 0 and zy0 == 0:
        return 0
    z = complex(zx0, zy0)
    logf = log(abs(z))
    fac = 1/2
    n = 0
    # if <= 0, always in the julia set
    if logf <= 0:
        return 0
    else:
        # implement the algorithm found in 2
        while n < 50 and abs(z)<1e24:
            z = z*z+c
            logf = logf +fac*log(abs(1+c/(z*z)))
            n = n+1
            fac = fac*fac
        return logf

# input z and c, output a 3 tuple of RGB value
def u(zx0, zy0, c):
    x = green_function(zx0, zy0, c)
    if x <= 0:
        return 0, 0, 0
    else:
        r = (128+floor(10000*x/128)) % 256
        g = floor(log2(log2(x))) % 256
        b = floor(100*log2(x)) % 256
        return r, g, b

# input z and c, output a 3 tuple of RGB value
def v(zx0, zy0, c):
    x = green_function(zx0, zy0, c)
    if x <= 1:
        return 0, 0, 0
    else:
        r = floor(255*(1+cos(2*sqrt(2)*pi*log(x)))/2)
        g = floor(255*(1+cos(2*sqrt(3)*pi*log(x)))/2)
        b = floor(255*(1+cos(2*sqrt(3)*pi*log(log(x))))/2)
        return r, g, b

# input the pix matrix and c. no output, just coloring by u
def julia_u(pix, c):
    global xmin, xmax, ymin, ymax, deltax, deltay, w, h
    cx = c.real
    cy = c.imag
    for i in range(0, w):
        for j in range(0, h):
            x = xmin + i * deltax
            y = ymin + j * deltay
            # init z0
            zx = x
            zy = y
            maxiter = 100
            snorm = x * x + y * y
            iterations = 0
            while (snorm <= 1000) and (iterations <= maxiter):
                a = zx * zx - zy * zy + cx
                b = 2 * zx * zy + cy
                zx = a
                zy = b
                snorm = zx * zx + zy * zy
                iterations += 1
            snorm = zx * zx + zy * zy
            if (snorm >= 1000):
                # not in the julia set, needs coloring
                pix[i, h - j - 1] = u(zx, zy, c)
            else:
                # in the julia set, just set it as black
                pix[i, h - j - 1] = (0, 0, 0)

# input the pix matrix and c. no output, just coloring by v
def julia_v(pix, c):
    global xmin, xmax, ymin, ymax, deltax, deltay, w, h, white, black
    cx = c.real
    cy = c.imag
    for i in range(0, w):
        for j in range(0, h):
            x = xmin + i * deltax
            y = ymin + j * deltay
            # init z0
            zx = x
            zy = y
            maxiter = 100
            snorm = x * x + y * y
            iterations = 0
            while (snorm <= 1000) and (iterations <= maxiter):
                a = zx * zx - zy * zy + cx
                b = 2 * zx * zy + cy
                zx = a
                zy = b
                snorm = zx * zx + zy * zy
                iterations += 1
            snorm = zx * zx + zy * zy
            if (snorm >= 1000):
                # not in the julia set
                pix[i, h - j - 1] = v(zx, zy, c)
            else:
                # in the julia set
                pix[i, h - j - 1] = (0, 0, 0)

# input the pix matrix no output, just coloring by u
def mandelbrot_u(pix):
    global xmin,xmax, ymin,ymax, deltax,deltay,w,h

    for i in range(0,w):
        for j in range (0,h):
            cx=xmin+ i *deltax
            cy=ymin+j*deltay
            c = complex(cx, cy)
            zx=0
            zy=0
            maxiter=100
            snorm= zx*zx+zy*zy
            iterations=0
            while ((snorm<=1000) and (iterations<= maxiter)):
                a= zx*zx - zy*zy + cx
                b= 2*zx* zy + cy
                zx=a
                zy=b
                snorm= zx*zx + zy*zy
                iterations+=1
            snorm= zx*zx + zy*zy
            if (snorm >=1000):
                ### not in mandelbrot s
                pix[i,h-j-1] = u(zx, zy, c)
            else:
                pix[i,h-j-1]=(0, 0, 0)

# input the pix matrix no output, just coloring by v
def mandelbrot_v(pix):
    global xmin,xmax, ymin,ymax, deltax,deltay,w,h

    for i in range(0,w):
        for j in range (0,h):
            cx=xmin+ i *deltax
            cy=ymin+j*deltay
            c = complex(cx, cy)
            zx=0
            zy=0
            maxiter=100
            snorm= zx*zx+zy*zy
            iterations=0
            while ((snorm<=1000) and (iterations<= maxiter)):
                a= zx*zx - zy*zy + cx
                b= 2*zx* zy + cy
                zx=a
                zy=b
                snorm= zx*zx + zy*zy
                iterations+=1
            snorm= zx*zx + zy*zy
            if (snorm >=1000):
                ### not in mandelbrot s
                pix[i,h-j-1] = v(zx, zy, c)
            else:
                pix[i,h-j-1]=(0, 0, 0)

c0 = complex(0, 0)
c1 = complex(-0.123, 0.745)
c2 = complex(-1, 0)
c3 = complex(0.3, 0.5)
c4 = complex(-0.8, 0.156)
C = c0, c1, c2, c3, c4
# Exercice 1 (3)

for c in C:
    i = 0
    julia_u(pix, c)
    bitmap.save("julia_u"+str(c)+".jpg")

# Exercise 1 (4)

for c in C:
    julia_v(pix, c)
    bitmap.save("julia_v"+str(c)+".jpg")

# Exercise 2 (3)

mandelbrot_u(pix)
bitmap.save("mandelbrot_u.jpg")

# Exercise 2 (4)

mandelbrot_v(pix)
bitmap.save("mandelbrot_v.jpg")

# Exercise 2 (5)

xmin1 = -0.25
xmax1 = 0.75
ymin1 = -0.5
ymax1 = 0.5
# umbers columns is 500
w1 = 500
h1 = 500
bitmap1 = Image.new("RGB", (w1, h1), "white")
pix1 = bitmap1.load()
deltax1 = abs(xmax1 - xmin1)/w1
deltay1 = abs(ymax1 - ymin1)/h1

# not going to be used, just another coloring function
def bdy(zx0, zy0, c):
    x = green_function(zx0, zy0, c)
    if x <= 0:
        return 0, 0, 0
    else:
        r = floor(255*(1+cos(2*pi*x))/2)
        g = floor(255*(1+cos(2*pi*x))/2)
        b = floor(255*(1+cos(2*pi*x))/2)
        return r, g, b

# coloring function
def bdy1(zx0, zy0, c):
    x = green_function(zx0, zy0, c)
    if x <= 0:
        return 0, 0, 0
    else:
        b = 7
        r = floor(255*(1+cos(2*sqrt(7)*pi*log(x)))/2)
        g = floor(255*(1-cos(b*x))/2)
        b = floor(255*(1+cos(2*sqrt(5)*pi*log(log(x))))/2)
        return r, g, b

# input the pix matrix no output, just coloring by bdy1
def mandelbrot_bdy1(pix):
    global xmin1,xmax1, ymin1,ymax1, deltax1,deltay1,w1,h1

    for i in range(0,w1):
        for j in range (0,h1):
            cx=xmin1 + i * deltax1
            cy=ymin1 + j * deltay1
            c = complex(cx, cy)
            zx=0
            zy=0
            maxiter=100
            snorm= zx*zx+zy*zy
            iterations=0
            while ((snorm<=1000) and (iterations<= maxiter)):
                a= zx*zx - zy*zy + cx
                b= 2*zx* zy + cy
                zx=a
                zy=b
                snorm= zx*zx + zy*zy
                iterations+=1
            snorm= zx*zx + zy*zy
            if (snorm >=1000):
                # not in mandelbrot s, needs coloring
                pix[i,h-j-1] = bdy1(zx, zy, c)
            else:
                pix[i,h-j-1]=(0, 0, 0)

mandelbrot_bdy1(pix1)
bitmap1.save("mandelbrot_bdy1.jpg")

Project 3 Coloring Julia Set/project_3_gao.py


from django.shortcuts import render

from .forms import Contactforms
# Create your views here.

def index(request):
  if request.method == "POST" :
    contact = Contactforms(request.POST)
    if contact.is_valid:
       contact.save()

  context = {
     'Contactforms': Contactforms(),
  }

  return render(request, 'pages/index.html' , context)


import time
import torch.nn.functional as func
import torch.nn.functional
from comm.comm_utils import *
from modules.generation_utils import get_logits_warper
import json
from typing import List, Dict
from coordinator.coordinator_client import LocalCoordinatorClient

class DistInferenceMaskTokenPipeAutoBatch:
    def __init__(self, args, device, coord_client: LocalCoordinatorClient = None):
        print("=======Initialize Dist Inference(DistInferenceMaskTokenPipeAutoBatch).=======")
        self.has_work = torch.zeros(1, dtype=torch.int, device=device)
        self.task_settings = []
        self.dist_store = dist.distributed_c10d._get_default_store()
        self.task_count = 0
        self.current_job_ids = []
        self.device = device
        if coord_client is not None:
            self.coord_client = coord_client
        else:
            self.coord_client = None  # get_coordinator_client()
        self.init_job_id = args.job_id  # This should be changed to model ID later.
        # Model info for pipeline
        self.max_layers = args.max_layers
        self.num_layers = args.num_layers
        self.model_name = args.model_name
        self.model_type = args.model_type
        self.embedding_dim = self._get_embedding_size()
        self._layer_begin = args.num_layers * get_pipeline_parallel_rank()
        self._layer_end = args.num_layers * (get_pipeline_parallel_rank() + 1)
        if args.fp16:
            self.use_fp16 = True
            print("=======Gpipe use FP16=======")
        else:
            self.use_fp16 = False
            print("=======Gpipe use FP32=======")
        self.dtype = torch.float16 if self.use_fp16 else torch.float32
        from task_datasets.inference_data import get_tokenizer
        self.tokenizer = get_tokenizer(args)

        # Inference setting hyper-parameters
        self.top_k = args.top_k  # This variable cannot be updated interactively yet.
        self.echo_prompt = []
        self.num_completions = []
        self.top_k_per_token = []
        self.input_seq_length = []
        self.generate_seq_length = []
        self.logits_warpers = []
        self.i_current_token = []
        # self._init_batch_settings(task_settings)

        # Pipeline info
        self.pipeline_group_size = args.pipeline_group_size
        self.pp_rank = get_pipeline_parallel_rank()  # Rank is the pipeline rank by default.
        self.pre_node_rank = self.pp_rank - 1
        self.post_node_rank = self.pp_rank + 1 if self.pp_rank != self.pipeline_group_size - 1 else -1
        self.comm = get_pipeline_parallel_comm()
        self.seq_micro_batch_size = 1
        self.token_micro_batch_size = 1

        self.layers = {}
        self._create_layers()

        if self.pp_rank == self.pipeline_group_size - 1:
            self.ret_tokens = []
            self.ret_token_logprobs = []
            self.ret_topk_tokens = []
            self.ret_topk_token_logprobs = []
        if self.pp_rank == 0:
            self.recv_new_token = []
        if self.pp_rank == self.pipeline_group_size - 1:
            self.send_new_tokens = []
        self.input_seq_emb = []
        self.output_seq_emb = []
        self.input_token_emb = []
        self.output_token_emb = []
        # self._init_buffers()
        self.cached_attention = []
        # self._init_cached_seqs_and_attentions()

    def sync_has_work(self):
        self.comm.broadcast(self.has_work, src=0)
        print("sync_has_work:", self.has_work.item())
        self.comm.barrier()

    def _init_batch_settings(self):
        self.echo_prompt.clear()
        self.num_completions.clear()
        self.top_k_per_token.clear()
        self.input_seq_length.clear()
        self.generate_seq_length.clear()
        self.logits_warpers.clear()

        self.seq_num = len(self.task_settings)
        for i in range(self.seq_num):
            self.echo_prompt.append(self.task_settings[i].get('echo', False))
            self.num_completions.append(self.task_settings[i].get('n', 1))
            self.top_k_per_token.append(self.task_settings[i].get('logprobs', 0))
            current_seq_length = self.tokenizer(self.task_settings[i]['prompt'], return_tensors='pt', padding=True,
                                                truncation=False)['input_ids'].size(1)
            # 2048 is hardcoded.
            current_seq_length = min(current_seq_length, 2048 - self.task_settings[i].get('max_tokens', 1))
            self.input_seq_length.append(current_seq_length)
            self.generate_seq_length.append(self.task_settings[i].get('max_tokens', 1))

            current_top_p = self.task_settings[i].get('top_p', 1.0)
            current_temperature = self.task_settings[i].get('temperature', 0)
            current_logits_warper = get_logits_warper(
                top_k=(None if self.top_k is None or self.top_k == 0 else self.top_k),
                top_p=(None if current_top_p is None or current_top_p <= 0 else current_top_p),
                temperature=current_temperature,
                num_beams=1,
            )
            self.logits_warpers.append(current_logits_warper)
            self.i_current_token.append(None)

    def _print_batch_settings(self):
        print("<DistInferenceMaskTokenPipeAutoBatch-_print_batch_settings>:")
        for i in range(self.seq_num):
            print("-----------------------------------------------------------")
            print(f"Seq index <{i}>")
            print(f"echo_prompt: <{self.echo_prompt[i]}>")
            print(f"num_completions: <{self.num_completions[i]}>")
            print(f"top_k_per_token: <{self.top_k_per_token[i]}>")
            print(f"input_seq_length: <{self.input_seq_length[i]}>")
            print(f"generate_seq_length: <{self.generate_seq_length[i]}>")

    def _create_layers(self):
        if self.model_type == 'gpt2':
            from modules.hf_gpt2_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'gptj':
            from modules.hf_gptj_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'gptneox':
            from modules.hf_gptneox_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'opt':
            from modules.hf_opt_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'bloom':
            from modules.hf_bloom_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'yalm':
            from modules.yalm_module import GPTEmbeddings, GPTBlock, GPTLMHead
        elif self.model_type == 'glm':
            from modules.glm_module import GPTEmbeddings, GPTBlock, GPTLMHead
        else:
            raise Exception(f'unknown model type {self.model_type}')

        if self.pp_rank == 0:
            self.layers['emb'] = GPTEmbeddings.from_pretrained(self.model_name).to(self.dtype).eval().to(self.device)
        for layer_index in range(self.num_layers):
            # global layer indexing could be an argument
            global_layer_index = self.num_layers * self.pp_rank + layer_index
            # in case the total number of layers are not dividable by pipeline group size.
            if self.max_layers is not None and global_layer_index >= self.max_layers:
                self.num_layers = layer_index
                break
            print(f'loading layer {global_layer_index}')
            self.layers['block' + str(layer_index)] = GPTBlock.from_pretrained(
                self.model_name, layer_index=global_layer_index).to(self.dtype).eval().to(self.device)
            if self.coord_client:
                self.coord_client.update_status(self.init_job_id, 'running', returned_payload={
                    'rank': self.pp_rank, 'loaded_layer': layer_index, 'total_layer': self.num_layers})
        if self.pp_rank == self.pipeline_group_size - 1:
            self.layers['lm'] = GPTLMHead.from_pretrained(
                self.model_name
            ).to(self.dtype).eval().to(self.device)

    def _get_embedding_size(self):
        if self.model_type == 'gpt2':
            from modules.hf_gpt2_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.n_embd
        elif self.model_type == 'gptj':
            from modules.hf_gptj_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.n_embd
        elif self.model_type == 'gptneox':
            from modules.hf_gptneox_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.hidden_size
        elif self.model_type == 'opt':
            from modules.hf_opt_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.hidden_size
        elif self.model_type == 'bloom':
            from modules.hf_bloom_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.hidden_size
        elif self.model_type == 'yalm':
            from modules.yalm_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.hidden_size
        elif self.model_type == 'glm':
            from modules.glm_module import GPTConfig
            config = GPTConfig.from_pretrained(self.model_name)
            return config.hidden_size
        else:
            raise Exception(f'unknown model type {self.model_type}')

    def _init_buffers(self):
        if self.pp_rank == self.pipeline_group_size - 1:
            self.ret_tokens.clear()
            self.ret_token_logprobs.clear()
            self.ret_topk_tokens.clear()
            self.ret_topk_token_logprobs.clear()
        if self.pp_rank == 0:
            self.recv_new_token.clear()
        if self.pp_rank == self.pipeline_group_size - 1:
            self.send_new_tokens.clear()
        self.input_seq_emb.clear()
        self.output_seq_emb.clear()
        self.input_token_emb.clear()
        self.output_token_emb.clear()

        if self.pp_rank == self.pipeline_group_size - 1:
            for i in range(self.seq_num):
                current_ret_seq_length = self.generate_seq_length[i] if not self.echo_prompt[i] else \
                    self.input_seq_length[i] + self.generate_seq_length[i]
                current_ret_tokens = torch.zeros((self.num_completions[i], current_ret_seq_length),
                                                 requires_grad=False, device=self.device, dtype=torch.int64)
                self.ret_tokens.append(current_ret_tokens)
                current_ret_token_logprobs = torch.zeros((self.num_completions[i], current_ret_seq_length),
                                                         requires_grad=False, device=self.device, dtype=self.dtype)
                self.ret_token_logprobs.append(current_ret_token_logprobs)
                if self.top_k_per_token[i] > 0:
                    current_ret_topk_tokens = torch.zeros((self.num_completions[i], current_ret_seq_length,
                                                           self.top_k_per_token[i]), requires_grad=False,
                                                          device=self.device, dtype=torch.int64)
                    self.ret_topk_tokens.append(current_ret_topk_tokens)
                    current_ret_topk_token_logprobs = torch.zeros((self.num_completions[i], current_ret_seq_length,
                                                                   self.top_k_per_token[i]), requires_grad=False,
                                                                  device=self.device, dtype=self.dtype)
                    self.ret_topk_token_logprobs.append(current_ret_topk_token_logprobs)
                else:
                    self.ret_topk_tokens.append(None)
                    self.ret_topk_token_logprobs.append(None)
        if self.pp_rank == 0:
            for i in range(self.seq_num):
                current_recv_new_token = torch.zeros((self.num_completions[i], 1), requires_grad=False,
                                                     device=self.device, dtype=torch.int64)
                self.recv_new_token.append(current_recv_new_token)
        if self.pp_rank == self.pipeline_group_size - 1:
            for i in range(self.seq_num):
                current_send_new_tokens = torch.zeros((self.num_completions[i], 1), requires_grad=False,
                                                      device=self.device, dtype=torch.int64)
                self.send_new_tokens.append(current_send_new_tokens)
        for i in range(self.seq_num):
            if self.generate_seq_length[i] == 0:
                current_input_seq_emb = torch.zeros((1, self.input_seq_length[i] - 1, self.embedding_dim),
                                                    requires_grad=False, device=self.device, dtype=self.dtype)
                current_output_seq_emb = torch.zeros((1, self.input_seq_length[i] - 1, self.embedding_dim),
                                                     requires_grad=False, device=self.device, dtype=self.dtype)
            else:
                current_input_seq_emb = torch.zeros((1, self.input_seq_length[i], self.embedding_dim),
                                                    requires_grad=False, device=self.device, dtype=self.dtype)
                current_output_seq_emb = torch.zeros((1, self.input_seq_length[i], self.embedding_dim),
                                                     requires_grad=False, device=self.device, dtype=self.dtype)
            self.input_seq_emb.append(current_input_seq_emb)
            self.output_seq_emb.append(current_output_seq_emb)
            current_input_token_emb = torch.zeros((self.num_completions[i], 1, self.embedding_dim),
                                                  requires_grad=False, device=self.device, dtype=self.dtype)
            self.input_token_emb.append(current_input_token_emb)
            current_output_token_emb = torch.zeros((self.num_completions[i], 1, self.embedding_dim),
                                                   requires_grad=False, device=self.device, dtype=self.dtype)
            self.output_token_emb.append(current_output_token_emb)

    def _print_buffers(self):
        print(f"<DistInferenceMaskTokenPipeAutoBatch-_print_buffers>: rank-<{self.pp_rank}>=======")
        for i in range(self.seq_num):
            print("-----------------------------------------------------------")
            print(f"Seq index <{i}>")
            if self.pp_rank == self.pipeline_group_size - 1:
                print(f"ret_tokens: <{self.ret_tokens[i].shape}>")
                print(f"ret_token_logprobs: <{self.ret_token_logprobs[i].shape}>")
                if self.ret_topk_tokens[i] is not None:
                    print(f"ret_topk_tokens: <{self.ret_topk_tokens[i].shape}>")
                else:
                    print(f"ret_topk_tokens: <None>")
                if self.ret_topk_token_logprobs[i] is not None:
                    print(f"ret_topk_token_logprobs: <{self.ret_topk_token_logprobs[i].shape}>")
                else:
                    print(f"ret_topk_token_logprobs: <None>")
            if self.pp_rank == 0:
                print(f"recv_new_token: <{self.recv_new_token[i].shape}>")
            if self.pp_rank == self.pipeline_group_size - 1:
                print(f"send_new_tokens: <{self.send_new_tokens[i].shape}>")
            print(f"input_seq_emb: <{self.input_seq_emb[i].shape}>")
            print(f"output_seq_emb: <{self.output_seq_emb[i].shape}>")
            print(f"input_token_emb: <{self.input_token_emb[i].shape}>")
            print(f"output_token_emb: <{self.output_token_emb[i].shape}>")

    def update_batch_setting(self, task_settings: List[Dict] = None, job_ids: List[str] = None):
        print(f"<DistInferenceMaskTokenPipeAutoBatch-update_batch_setting>: rank-<{self.pp_rank}>=======")
        print(task_settings)
        print(job_ids)
        self.task_count += 1
        if self.pp_rank == 0:
            assert task_settings is not None and job_ids is not None
            task_settings_str = json.dumps({'task_settings': task_settings, 'job_ids': job_ids})
            self.dist_store.set('current_task_'+str(self.task_count), task_settings_str)
            print(task_settings_str)
            self.task_settings.clear()
            self.task_settings.extend(task_settings)
            self.current_job_ids.clear()
            self.current_job_ids.extend(job_ids)
        else:
            assert task_settings is None
            task_settings_str = self.dist_store.get('current_task_'+str(self.task_count))
            print(task_settings_str)
            task_dict = json.loads(task_settings_str)
            task_settings = task_dict['task_settings']
            self.task_settings.clear()
            self.task_settings.extend(task_settings)
            self.current_job_ids.clear()
            self.current_job_ids.extend(task_dict['job_ids'])
        self._init_batch_settings()
        self._print_batch_settings()
        self._init_buffers()
        self._print_buffers()
        self._init_cached_seqs_and_attentions()

    def _init_cached_seqs_and_attentions(self):
        for i in range(self.seq_num):
            if not self.echo_prompt[i]:
                self.i_current_token[i] = 0
            else:
                self.i_current_token[i] = self.input_seq_length[i]
        self.cached_attention.clear()
        for _ in range(self.num_layers):
            self.cached_attention.append([None for _ in range(self.seq_num)])
        self.token_cached_attention = []
        for _ in range(self.num_layers):
            self.token_cached_attention.append([None for _ in range(self.seq_num)])

    def _generate_new_token(self, index):
        assert self.pp_rank == self.pipeline_group_size - 1
        z = self.layers['lm'](self.output_token_emb[index])
        if torch.isnan(z).any():
            print('containing nan, setting them to zero!')
            print(z)
        z = z.float().nan_to_num()  # test if fp32 whether cause inf
        z = torch.nn.functional.log_softmax(z, -1)

        if self.top_k_per_token[index] > 0:
            logprobs, indices = z.topk(k=self.top_k_per_token[index], dim=-1)
            self.ret_topk_tokens[index][:, self.i_current_token[index]] = indices.squeeze(1)
            self.ret_topk_token_logprobs[index][:, self.i_current_token[index]] = logprobs.squeeze(1)
        # [:, -1] because multinomial only accept 1/2d tensors
        z_to_sample = z[:, -1]  # bs, vocab
        z_to_sample = self.logits_warpers[index](None, z_to_sample)
        p_to_sample = z_to_sample.softmax(-1).clamp(0, 1).nan_to_num()
        indices = torch.multinomial(p_to_sample, num_samples=1)  # bs, 1
        logprobs = torch.gather(z[:, -1], -1, indices)  # bs, 1
        self.send_new_tokens[index] = indices

        self.ret_tokens[index][:, self.i_current_token[index]] = indices.squeeze(-1)
        self.ret_token_logprobs[index][:, self.i_current_token[index]] = logprobs.squeeze(-1)
        self.i_current_token[index] += 1

    def _merge_cached_seqs_and_attentions(self):
        for i in range(self.seq_num):
            if not self.echo_prompt[i]:
                self.i_current_token[i] = 0
            else:
                self.i_current_token[i] = self.input_seq_length[i]
        if self.pp_rank == self.pipeline_group_size - 1:
            for i in range(self.seq_num):
                if self.generate_seq_length[i] != 0:
                    self._generate_new_token(i)

    def _copy_initial_token_emb(self, index):
        assert self.pp_rank == self.pipeline_group_size - 1
        print("_copy_initial_token_emb")
        for k in range(self.num_completions[index]):
            print(f"_copy_initial_token_emb {k}/{self.num_completions[index]}")
            self.output_token_emb[index][k] = self.output_seq_emb[index][:, -1:]

    def _generate_echo_token_logprobs(self, index, indices):
        assert self.pp_rank == self.pipeline_group_size - 1
        assert self.num_completions[index] == 1

        if self.generate_seq_length[index] == 0:
            z = self.layers['lm'](self.output_seq_emb[index])
        else:
            z = self.layers['lm'](self.output_seq_emb[index][:, :-1])
        z = func.log_softmax(z, -1)
        original_indices = indices
        indices = indices[:, 1:]  # skip first

        logprobs = torch.gather(z, -1, indices.unsqueeze(-1)).squeeze(-1)
        self.ret_tokens[index][:, :self.i_current_token[index]] = original_indices
        self.ret_token_logprobs[index][:, 1:self.i_current_token[index]] = logprobs
        if self.top_k_per_token[index] > 0:
            logprobs, indices = z.topk(k=self.top_k_per_token[index], dim=-1)
            self.ret_topk_tokens[index][:, 1:self.i_current_token[index]] = indices
            self.ret_topk_token_logprobs[index][:, 1:self.i_current_token[index]] = logprobs

    def _forward_compute_prompt_seq(self, index, seq, mask):
        print("Compute prompt seq<", index, ">.")
        if self.pp_rank == 0:
            self.input_seq_emb[index] = self.layers['emb'](seq, mask=mask)
        current_emb = self.input_seq_emb[index]
        caches = [None] * self.num_layers
        for layer_index in range(self.num_layers):
            current_emb, caches[layer_index] = \
                self.layers['block' + str(layer_index)](current_emb, caches[layer_index], mask=mask)
            self.cached_attention[layer_index][index] = caches[layer_index]
        self.output_seq_emb[index] = current_emb
        if self.pp_rank == self.pipeline_group_size - 1:
            self._copy_initial_token_emb(index)
            if self.echo_prompt[index]:
                self._generate_echo_token_logprobs(index, indices=seq)

    def _get_cached_attention(self, layer_index, index):
        if self.num_completions[index] == 1:
            return self.cached_attention[layer_index][index]
        else:
            # 2* (token_bs, ., seq_len, .)
            prompt_cache = self.cached_attention[layer_index][index]
            # 2* (token_bs * num_completion, ., seq_len, .)
            token_cache = self.token_cached_attention[layer_index][index]
            # 2*(token_bs * num_completion, ., seq_len, .)
            prompt_cache = [prompt_cache[0].repeat(self.num_completions[index], 1, 1, 1),
                            prompt_cache[1].repeat(self.num_completions[index], 1, 1, 1)]
            if token_cache is not None:
                token_cache = [torch.cat([prompt_cache[0], token_cache[0]], dim=2),
                               torch.cat([prompt_cache[1], token_cache[1]], dim=2)]
            else:
                token_cache = prompt_cache
            return token_cache

    def _set_cached_attention(self, cache, layer_index, index):
        if self.num_completions[index] == 1:
            self.cached_attention[layer_index][index] = cache
        else:
            self.token_cached_attention[layer_index][index] = [cache[0][:, :, self.input_seq_length[index]:],
                                                               cache[1][:, :, self.input_seq_length[index]:]]

    def _forward_compute_generate_token(self, index, mask=None):
        if mask is not None and self.num_completions[index] > 1:
            # repeat n times
            mask = mask.repeat(self.num_completions[index], 1)
        # print("Compute generate seq micro-batch <", index, ">.")
        if self.pp_rank == 0:
            current_emb = self.layers['emb'](self.recv_new_token[index], self.cached_attention[0][index], mask=mask)
        else:
            current_emb = self.input_token_emb[index]
        for layer_index in range(self.num_layers):
            cache = self._get_cached_attention(layer_index, index)
            current_emb, cache = self.layers['block' + str(layer_index)](current_emb, cache, mask=mask)
            self._set_cached_attention(cache, layer_index, index)
        self.output_token_emb[index] = current_emb
        if self.pp_rank == self.pipeline_group_size - 1:
            self._generate_new_token(index)

    def _process_mask_during_generation(self, attention_mask):
        for i in range(self.seq_num):
            if attention_mask[i] is not None:
                # increase one for the new token
                attention_mask[i] = func.pad(attention_mask[i], pad=(0, 1), mode='constant', value=1)
        return attention_mask

    def forward_seq_pipeline_stage(self, input_data=None, attention_mask=None):
        # if self.pp_rank == 0 or self.pp_rank == self.pipeline_group_size - 1:
        if self.pp_rank == 0:
            assert (input_data is not None)
            if self.pp_rank == 0:
                for i in range(self.seq_num):
                    if self.generate_seq_length[i] == 0:
                        # input reduce 1 for first node
                        input_data[i] = input_data[i][:-1]
        if input_data is not None:
            input_seqs = input_data
        else:
            input_seqs = [None] * self.seq_num

        if attention_mask is not None:
            for i in range(self.seq_num):
                if attention_mask[i] is not None:
                    if self.generate_seq_length[i] == 0:
                        # attention reduce 1
                        attention_mask[i] = attention_mask[i][:-1]
        else:
            attention_mask = [None] * self.seq_num

        for i in range(self.seq_num):
            if self.pp_rank == 0:  # Only send output to next node, do not receive
                # Compute
                self._forward_compute_prompt_seq(index=i, seq=input_seqs[i], mask=attention_mask[i])
                # Send
                self.comm.send(self.output_seq_emb[i], dst=self.post_node_rank)
            elif self.pp_rank == self.pipeline_group_size - 1:  # Only receive input from last node, do not send
                # Receive
                self.comm.recv(self.input_seq_emb[i], src=self.pre_node_rank)
                # Compute
                self._forward_compute_prompt_seq(index=i, seq=input_seqs[i], mask=attention_mask[i])
            else:  # receive, compute, and send
                # Receive
                self.comm.recv(self.input_seq_emb[i], src=self.pre_node_rank)
                # Compute
                self._forward_compute_prompt_seq(index=i, seq=input_seqs[i], mask=attention_mask[i])
                # Send
                self.comm.send(self.output_seq_emb[i], dst=self.post_node_rank)

    def forward_new_token_pipeline_stage(self, attention_mask=None):
        self._merge_cached_seqs_and_attentions()
        for step in range(max(self.generate_seq_length)):
            print("Compute generate token step <", step, ">.")
            # for last node, the first step does not compute
            if step != 0 or self.pp_rank != self.pipeline_group_size - 1:
                if attention_mask:
                    attention_mask = self._process_mask_during_generation(attention_mask)
            self.forward_new_token_pipeline_step(step, attention_masks=attention_mask)

    def forward_new_token_pipeline_step(self, step: int, attention_masks=None):
        for i in range(self.seq_num):
            # This seq does not need these much seq.
            if step >= self.generate_seq_length[i]:
                continue

            # Last node:
            if self.pp_rank == self.pipeline_group_size - 1:
                if step == 0:
                    # Send
                    self.comm.send(self.send_new_tokens[i], dst=0)
                else:
                    # Receive
                    self.comm.recv(self.input_token_emb[i], src=self.pre_node_rank)
                    # Compute
                    self._forward_compute_generate_token(i, mask=attention_masks[i] if attention_masks else None)
                    if step != self.generate_seq_length[i] - 1:
                        # Send
                        self.comm.send(self.send_new_tokens[i], dst=0)
            # Rank-0 node:
            elif self.pp_rank == 0:
                if step != self.generate_seq_length[i] - 1:
                    # Receive
                    self.comm.recv(self.recv_new_token[i], src=self.pipeline_group_size - 1)
                    # Compute
                    self._forward_compute_generate_token(i, mask=attention_masks[i] if attention_masks else None)
                    # Send
                    self.comm.send(self.output_token_emb[i], dst=self.post_node_rank)
            else:  # Middle nodes:
                if step != self.generate_seq_length[i] - 1:
                    # Receive
                    self.comm.recv(self.input_token_emb[i], src=self.pre_node_rank)
                    # Compute
                    self._forward_compute_generate_token(i, mask=attention_masks[i] if attention_masks else None)
                    # Send
                    self.comm.send(self.output_token_emb[i], dst=self.post_node_rank)

    def inference_batch(self, input_=None, output_=None, attention_mask=None):
        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> Enter!")
        self.comm.barrier()
        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> after first barrier!")
        self._init_cached_seqs_and_attentions()  # TODO: should I put here
        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> after first _init_cached_seqs_and_attentions!")
        self.comm.barrier()
        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> after second barrier!")
        start_time = time.time()

        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> enter computation!")
        with torch.no_grad():
            self.forward_seq_pipeline_stage(input_data=input_, attention_mask=attention_mask)
            print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> forward_seq_pipeline_stage is done!")
            self.forward_new_token_pipeline_stage(attention_mask=attention_mask)
            print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> forward_new_token_pipeline_stage is done!")

        self.comm.barrier()
        print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> after third barrier!")

        if self.pp_rank == self.pipeline_group_size - 1 and output_ is not None:
            ret_tokens = []
            ret_token_logprobs = []
            ret_topk_tokens = []
            ret_topk_token_logprobs = []
            for i in range(self.seq_num):
                ret_tokens.append(self.ret_tokens[i][:, :self.i_current_token[i]].cpu())
                ret_token_logprobs.append(self.ret_token_logprobs[i][:, :self.i_current_token[i]].cpu())
                if self.top_k_per_token[i] > 0:
                    ret_topk_tokens.append(self.ret_topk_tokens[i][:, :self.i_current_token[i]].cpu())
                    ret_topk_token_logprobs.append(self.ret_topk_token_logprobs[i][:, :self.i_current_token[i]].cpu())
                else:
                    ret_topk_tokens.append(None)
                    ret_topk_token_logprobs.append(None)
            print(f"<AutoBatch inference_batch> rank-<{self.pp_rank}> after marker1 !")

            for i in range(self.seq_num):
                item = {
                    'token_ids': ret_tokens[i],
                    'token_logprobs': ret_token_logprobs[i],
                    'topk_ids': ret_topk_tokens[i],
                    'topk_logprobs': ret_topk_token_logprobs[i]
                }
                output_.append(item)
            print(f"<inference_batch> rank-<{self.pp_rank}> after marker2 !")

        end_time = time.time()
        iter_time = end_time - start_time
        print("<AutoBatch inference_batch> rank-{} INFERENCE iteration takes {:3.2f}s".format(self.pp_rank, iter_time))
        print("-------------------------------------------")
        return iter_time


# A part of NonVisual Desktop Access (NVDA)
# Copyright (C) 2006-2021 NV Access Limited, Babbage B.V., Joseph Lee, Mesar Hameed, kvark128
# This file is covered by the GNU General Public License.
# See the file COPYING for more details.

from collections import namedtuple
import math
import colorsys
from ctypes.wintypes import COLORREF
import re
from functools import lru_cache
from typing import Union

#: Flag to indicate color being decoded from displayModelFormatColor_t
# is transparent.
# See (displayModel.cpp in nvdaHelper) displayModelFormatColor_t::TRANSPARENT_BIT
TRANSPARENT_BITFLAG = 0x01 << 24
ALPHA_OPAQUE: int = 0xFF  # no transparency.
ALPHA_TRANSPARENT: int = 0x00  # completely transparent

class RGB(namedtuple('RGB',('red','green','blue'))):
	"""Represents a color as an RGB (red green blue) value"""

	#: The transparency of the color.
	# Note: may be useful for add-ons or app modules who wish to customize color reporting
	# of background colors of applications that rely on the display model.
	# These applications do not always reliably report their background color.
	# The background color may be reported for text with transparent backgrounds.
	# In other cases the background for the text is transparent, but the background color is
	# visually correct and rendered in a different way EG via GDI filledRect.
	alphaValue: int = ALPHA_OPAQUE  # no transparency by default

	@classmethod
	def fromDisplayModelFormatColor_t(cls, c: int) -> "RGB":
		"""factory method to create an RGB from a DisplayModelFormatColor_t
			Color format is 4 bytes:  0xTTbbggrr
			TT bit flags, only bit 1 used: set for transparent.
			Encoding alpha in TT was considered (eg 0xFFbbggrr for opaque),
			but this would break compatibility with code that continues to pass
			0x00bbggrr for opaque, all usages would need to be fixed.
		"""
		rr = c & 0xFF
		gg = (c >> 8) & 0xFF
		bb = (c >> 16) & 0xFF
		tt = c & TRANSPARENT_BITFLAG
		rgb = cls(rr, gg, bb)
		rgb.alphaValue = ALPHA_TRANSPARENT if bool(tt) else ALPHA_OPAQUE
		return rgb

	@classmethod
	def fromCOLORREF(cls, c: Union[COLORREF, int]) -> "RGB":
		"""factory method to create an RGB from a COLORREF ctypes instance
		COLORREF format is 4 bytes: 0x00bbggrr
		According to MSDN, COLORREF high order byte must be zero.
		Handling of int is kept to maintain backwards compatibility.
		"""
		if isinstance(c, COLORREF):
			c = c.value
		elif not isinstance(c, int):
			raise TypeError(c)
		return cls.fromDisplayModelFormatColor_t(c)

	_re_RGBFunctionString=re.compile(r'rgb\(\s*(\d+%?)\s*,\s*(\d+%?)\s*,\s*(\d+%?)\s*\)',re.I)
	_re_RGBAFunctionString=re.compile(r'rgba\(\s*(\d+%?)\s*,\s*(\d+%?)\s*,\s*(\d+%?)\s*,\s*\d+(\.\d+)?\s*\)',re.I)

	@staticmethod
	def _RGBStringValToInt(s):
		val=int(round(int(s[:-1])*2.55)) if s.endswith('%') else int(s)
		if val<0 or val>255:
			raise ValueError("%s out of range"%val)
		return val

	@classmethod
	def fromString(cls,s):
		"""
		Factory method to create an RGB instance from a css RGB string representation.
		"""
		s=s.strip()
		#Try to match on the form RGB(x,y,z)
		m=cls._re_RGBFunctionString.match(s) or cls._re_RGBAFunctionString.match(s)
		if m:
			r=cls._RGBStringValToInt(m.group(1))
			g=cls._RGBStringValToInt(m.group(2))
			b=cls._RGBStringValToInt(m.group(3))
			return RGB(r,g,b)
		if s.startswith('#'):
			sLen=len(s)
			try:
				val=int(s[1:],16)
			except ValueError:
				val=None
			if val is not None:
				#Perhaps its a #aarrggbb or #rrggbb hex string
				if sLen==7 or sLen==9:
					r=(val>>16)&0xff
					g=(val>>8)&0xff
					b=val&0xff
					return RGB(r,g,b)
				#Perhaps its a #argb or #rgb hex string
				if sLen==4 or sLen==5:
					r=((val>>8)&0xf)+(((val>>8)&0xf)<<4)
					g=((val>>4)&0xf)+(((val>>4)&0xf)<<4)
					b=(val&0xf)+((val&0xf)<<4)
					return RGB(r,g,b)
		raise ValueError("invalid RGB string: %s"%s)

	def toCOLORREF(self) -> COLORREF:
		"""Returns a COLORREF ctypes instance
		"""
		return COLORREF(self.red & 0xff | ((self.green & 0xff) << 8) | ((self.blue & 0xff) << 16))

	def toGDIPlusARGB(self, alpha: int = 255) -> int:
		"""Creates a GDI+ compatible ARGB color, using the specified alpha for the alpha component.
		@param alpha: The alpha part of the ARGB color,
			0 is fully transparent and 255 is fully opaque.
			Defaults to 255 (opaque).
		@type alpha: int
		"""
		return (alpha << 24) | (self.red << 16) | (self.green << 8) | self.blue

	@property
	def name(self):
		import config
		shouldReportTransparent = config.conf["documentFormatting"]["reportTransparentColor"]
		return _calcColorName(self.red, self.green, self.blue, self.alphaValue, shouldReportTransparent)

@lru_cache(maxsize=5000)
def _calcColorName(red: int, green: int, blue: int, alpha: int, reportTransparent: bool):
	# convert to hsv (hue, saturation, value)
	h, s, v = colorsys.rgb_to_hsv(red / 255.0, green / 255.0, blue / 255.0)
	sv = s * v
	if sv < 0.02:
		# There is not enough saturation to perceive a hue, therefore its on the scale from black to white.
		v = v * 100
		nv = min(shadeNames, key=lambda i: abs(i - v))
		closestName = shadeNames[nv]
	else:
		h *= 360
		s *= 100
		v *= 100
		# Find the closest named hue (red, orange, yellow...)
		nh = min(hueNames, key=lambda i: 180 - abs(abs(i - h) - 180))
		hueName = hueNames[nh]
		ns = min(brightnessLabelsBySaturation, key=lambda i: abs(i - s))
		brightnessLabels = brightnessLabelsBySaturation[ns]
		nv = min(brightnessLabels, key=lambda i: abs(i - v))
		if nh in brownHueNames:
			brownNV = brownBrightnessRedirectionValues.get(nv)
			if brownNV is not None:
				nv = brownNV
				hueName = brownHueNames[nh]
		variationTemplate = brightnessLabels.get(nv)
		if variationTemplate:
			closestName = variationTemplate.format(color=hueName)
		else:
			closestName = hueName
	# the color is transparent and expected to be reported.
	if alpha < ALPHA_OPAQUE and reportTransparent:
		closestName = pgettext(
			'color variation',
			# Translators: a transparent color, {colorDescription} replaced with the full description of the color e.g.
			# "transparent bright orange-yellow"
			'transparent {colorDescription}'
		).format(colorDescription=closestName)
	return closestName

shadeNames={
	# Translators: the color white (HSV saturation0%, value 100%)
	100:pgettext('color hue','white'),
	# Translators: the color very light grey (HSV saturation 0%, value 84%)
	84:pgettext('color hue','very light grey'),
	# Translators: the color light grey (HSV saturation 0%, value 66%)
	66:pgettext('color hue','light grey'),
	# Translators: the color grey (HSV saturation 0%, value 50%)
	50:pgettext('color hue','grey'),
	# Translators: the color dark grey (HSV saturation 0%, value 34%)
	34:pgettext('color hue','dark grey'),
	# Translators: the color very dark grey (HSV saturation 0%, value 16%)
	16:pgettext('color hue','very dark grey'),
	# Translators: the color black (HSV saturation 0%, value 0%)
	0:pgettext('color hue','black'),
}

hueNames={
	# Translators: The color red (HSV hue 0 degrees)
	0:pgettext('color hue','red'),
	# Translators: The color between  red and orange (HSV hue 15 degrees)
	15:pgettext('color hue','red-orange'),
	# Translators: The color orange (HSV hue 30 degrees)
	30:pgettext('color hue','orange'),
	# Translators: The color between  orange and yellow (HSV hue 45 degrees)
	45:pgettext('color hue','orange-yellow'),
	# Translators: The color yellow (HSV hue 60 degrees)
	60:pgettext('color hue','yellow'),
	# Translators: The color between  yellow and green (HSV hue 90 degrees)
	90:pgettext('color hue','yellow-green'),
	# Translators: The color green (HSV hue 120 degrees)
	120:pgettext('color hue','green'),
	# Translators: The color between  green and aqua (HSV hue 150 degrees)
	150:pgettext('color hue','green-aqua'),
	# Translators: The color aqua (HSV hue 180 degrees)
	180:pgettext('color hue','aqua'),
	# Translators: The color between  aqua and blue (HSV hue 210 degrees)
	210:pgettext('color hue','aqua-blue'),
	# Translators: The color blue (HSV hue 240 degrees)
	240:pgettext('color hue','blue'),
	# Translators: The color between  blue and purple (HSV hue 270 degrees)
	270:pgettext('color hue','blue-purple'),
	# Translators: The color purple (HSV hue 300 degrees)
	300:pgettext('color hue','purple'),
	# Translators: The color between  purple and pink (HSV hue 312 degrees)
	312:pgettext('color hue','purple-pink'),
	# Translators: The color pink (HSV hue 324 degrees)
	324:pgettext('color hue','pink'),
	# Translators: The color between  pink and red (HSV hue 342 degrees)
	342:pgettext('color hue','pink-red'),
}

brightnessLabelsBySaturation={
	100:{
		# Translators: a bright color (HSV saturation 100% and value 100%)
		100:pgettext('color variation','bright {color}'),
		# Translators: color (HSV saturation 100% and value 72%)
		72:pgettext('color variation','{color}'),
		# Translators: a dark color (HSV saturation 100% and value 44%)
		44:pgettext('color variation','dark {color}'),
		# Translators: a very dark color (HSV saturation 100% and value 16%)
		16:pgettext('color variation','very dark {color}'),
	},
	60:{
		# Translators: a light pale color (HSV saturation 50% and value 100%)
		100:pgettext('color variation','light pale {color}'),
		# Translators: a pale color (HSV saturation 50% and value 72%)
		72:pgettext('color variation','pale {color}'),
		# Translators: a dark pale color (HSV saturation 50% and value 44%)
		44:pgettext('color variation','dark pale {color}'),
		# Translators: a very dark color (HSV saturation 50% and value 16%)
		16:pgettext('color variation','very dark pale {color}'),
	},
	10:{
		# Translators: a light color almost white - hardly any hue (HSV saturation 10% and value 100%)
		100:pgettext('color variation','{color} white'),
		# Translators: a color almost grey - hardly any hue (HSV saturation 10% and value 72%)
		72:pgettext('color variation','{color} grey'),
		# Translators: a dark color almost grey - hardly any hue (HSV saturation 10% and value 44%)
		44:pgettext('color variation','dark {color} grey'),
		# Translators: a very dark color almost grey - hardly any hue (HSV saturation 10% and value 16%)
		16:pgettext('color variation','very dark {color} grey'),
	},
}

brownHueNames={
	# Translators: The color between  red and brown (HSV hue 15 degrees, below 50% brightness)
	15:pgettext('color hue','red-brown'),
	# Translators: The color brown (HSV hue 30 degrees, below 50% brightness)
	30:pgettext('color hue','brown'),
	# Translators: The color between  brown and yellow (HSV hue 45 degrees, below 50% brightness)
	45:pgettext('color hue','brown-yellow'),
}

brownBrightnessRedirectionValues={44:72,16:44}


"""
gptpose CLI
"""

import click
import os
import getpass

from gptpose.generator import PoseGenerator

def run(
    pose_description: str,
    output_path: str,
    model_name: str,
) -> None:
    """generates a pose image based on a body pose description."""

    generator = PoseGenerator(
        pose_description=pose_description, output_path=output_path
    )

    _ = generator.generate_pose(model_name=model_name)

@click.command()
@click.option(
    "-d",
    "--description",
    "pose_description",
    required=True,
    help="The body pose description",
)
@click.option(
    "-o",
    "--output-path",
    "output_path",
    required=True,
    help="The path to save the pose images.",
)
@click.option(
    "-n",
    "--model-name",
    "model_name",
    required=False,
    default="GPT-4",
    help="The OpenAI model name you want to use. https://platform.openai.com/docs/models",
)
def main(pose_description: str, output_path: str, model_name: str = "gpt-4"):
    if not os.environ.get("OPENAI_API_KEY"):
        os.environ["OPENAI_API_KEY"] = getpass.getpass(
            prompt="Enter your OpenAI API key: "
        )

    run(pose_description, output_path, model_name)

if __name__ == "__main__":
    main()


import pathlib

import pytest
import yaml

from zntrack import Node, Project, dvc, zn
from zntrack.utils import config

class SingleNode(Node):
    params_file = dvc.params()

    def run(self):
        pass

def test_dvc_params(proj_path):
    """Test serialized data as parameters"""
    params = {"a": 100}
    file = pathlib.Path("my_params.yaml")

    file.write_text(yaml.safe_dump(params))
    with Project() as proj:
        SingleNode(params_file=file)
    proj.run()

    assert SingleNode.from_rev().params_file == file
    dvc_file = yaml.safe_load(config.files.dvc.read_text())
    assert [{"my_params.yaml": None}] == dvc_file["stages"]["SingleNode"]["params"]

def test_dvc_params_list(proj_path):
    """Test serialized data as parameters"""
    params = {"a": 100}
    file1 = pathlib.Path("my_params1.yaml")
    file2 = pathlib.Path("my_params2.yaml")

    file1.write_text(yaml.safe_dump(params))
    file2.write_text(yaml.safe_dump(params))

    with Project() as proj:
        SingleNode(params_file=[file1, file2])
    proj.run()

    assert SingleNode.from_rev().params_file == [file1, file2]
    dvc_file = yaml.safe_load(config.files.dvc.read_text())
    assert [{"my_params1.yaml": None}, {"my_params2.yaml": None}] == dvc_file["stages"][
        "SingleNode"
    ]["params"]

class MixedParams(Node):
    params_file1 = dvc.params()
    params_file2 = dvc.params()
    params = zn.params()

    def run(self):
        pass

def test_dvc_mixed_params_list(proj_path):
    """Test serialized data as parameters"""
    params = {"a": 100}
    file1 = pathlib.Path("my_params1.yaml")
    file2 = pathlib.Path("my_params2.yaml")

    file1.write_text(yaml.safe_dump(params))
    file2.write_text(yaml.safe_dump(params))

    with Project() as proj:
        MixedParams(params_file1=[file1, file2], params_file2=file1, params={"a": 100})
    proj.run()

    assert MixedParams.from_rev().params_file1 == [file1, file2]
    assert MixedParams.from_rev().params_file2 == file1
    assert MixedParams.from_rev().params == {"a": 100}
    dvc_file = yaml.safe_load(config.files.dvc.read_text())
    assert [
        "MixedParams",
        {"my_params1.yaml": None},
        {"my_params2.yaml": None},
    ] == dvc_file["stages"]["MixedParams"]["params"]

def test_dvc_params_error(proj_path):
    """Test serialized data as parameters"""
    with pytest.raises(ValueError):
        with Project() as proj:
            SingleNode(params_file={"a": 100})
        proj.build()


import pymc3 as pm
import theano.tensor as tt
import pandas as pd
import pickle
import argparse
import sys

def parse_args(args):
    '''
    >>> parse_args(['-f', 'test.txt'])
    Namespace(f='test.txt', g='-1', s=None)
    '''

    parser = argparse.ArgumentParser(description="MCMC sampling of PHM.")
    parser.add_argument('-f', type=str,  metavar='filename',
                        default='/home/lee/cogsci/phm/oaksford_paper_data.csv',
                        help='Path to dataset for fitting')

    parser.add_argument('-n', type=int, metavar='n',
                        default=10, help='number of participants')

    return(parser.parse_args(args))

def get_posterior(data, n=100, draws=2000, n_init=200000, progressbar=True,
                  *args, **kwargs):
    with pm.Model() as model:

        # Define Priors
        p_err = pm.Uniform('p_err', 0, 0.1)  # Upper limit due to normalization
        p_ent = pm.Uniform('p_ent', 0, 1 - 6 * p_err)
        p_a = pm.Uniform('p_a', 0, 1 - 6 * p_err - p_ent)
        p_e = pm.Uniform('p_e', 0, 1 - 6 * p_err - p_ent)
        p_o = pm.Uniform('p_o', 0, 1 - 6 * p_err - p_ent)
        p_i = pm.Uniform('p_i', 0, 1 - 6 * p_err - p_ent)
        nvc_a = pm.Deterministic('nvc_a', 1 - p_a - 6 * p_err - p_ent)
        nvc_i = pm.Deterministic('nvc_i', 1 - p_i - 6 * p_err - p_ent)
        nvc_e = pm.Deterministic('nvc_e', 1 - p_e - 6 * p_err - p_ent)
        nvc_o = pm.Deterministic('nvc_o', 1 - p_o - 6 * p_err - p_ent)

        # Model specification: define all possible moods
        # syll tt-syntax a     i       e      o     NVC
        aa = [p_a,   p_ent, p_err, p_err, nvc_a]
        ai = [p_err, p_a,   p_err, p_ent, nvc_a]
        ia = ai
        ae = [p_err, p_err, p_a,   p_ent, nvc_a]
        ea = ae
        ao = [p_err, p_ent, p_err, p_a,   nvc_a]
        oa = ao
        ii = [p_err, p_i,   p_err, p_ent, nvc_i]
        ie = [p_err, p_err, p_i,   p_ent, nvc_i]
        ei = ie
        io = [p_err, p_ent, p_err, p_i, nvc_i]
        oi = io
        ee = [p_err, p_err, p_e,   p_ent, nvc_e]
        eo = [p_err, p_ent, p_err, p_e, nvc_e]
        oe = eo
        oo = [p_err, p_ent, p_err, p_o, nvc_o]

        # Define the relationship between moods and syllogisms
        moods = [aa, ai, ae, ao, ia, ii, ie, io, ea, ei, ee, eo, oa, oi, oe, oo]
        syllogs = []
        for m in moods:
            # Figure 1
            line = m[0:4] + [p_err] * 4 + [m[-1]]
            syllogs += [line]
            # Figure 2
            line = [p_err] * 4 + m[0:4] + [m[-1]]
            syllogs += [line]

            line = []
            for para in m[0:4]:
                if para == p_err:
                    line += [p_err]
                else:
                    line += [para / 2]
            # Paste this two times
            line *= 2
            # Add NVC
            line += [m[-1]]

            syllogs += [line] * 2

        model_matrix = tt.stack(syllogs)

        # Define likelihood
        pm.Multinomial(name='rates', n=n, p=model_matrix, observed=data)
        map_estimate = pm.find_MAP(model=model)

        trace = pm.sample(draws=draws, njobs=1, start=map_estimate,
                          n_init=n_init, progressbar=progressbar)

        print('Model logp = ', model.logp(map_estimate))
        return model, trace

def get_data(file_name, nr_part):
    df = pd.read_csv(file_name)
    # Rename columns
    df.columns = ('syllog', 'a', 'i', 'e', 'o', 'nvc')
    df.pop('syllog')
    df.astype(float)

    # Number of participants
    df = df.multiply(nr_part)
    df = df.round().astype(int)

    # Select the experimental data only !! NOT NORMALIZED, sum may be greater or less!!
    # exp_data = df[['a', 'i', 'e', 'o', 'nvc']]

    # Normalize experimental data via NVC
    data = df[['a', 'i', 'e', 'o']].copy()
    data['nvc'] = nr_part - data.sum(axis=1)

    return(data)

if __name__ == '__main__':
    # Read command line arguments
    args = parse_args(sys.argv[1:])

    # Parse and preprocess data
    data = get_data(file_name=args.f, nr_part=args.n)

    out_name = 'test'
    trace, model = get_posterior(data=data, n=args.n, draws=2000, n_init=20000)
    # Save thinned and burned trace
    # trace = trace[500, 5]
    with open('./' + out_name + '.pickle', 'wb') as pickle_file:
        pickle.dump({'model': model, 'trace': trace}, pickle_file,
                    pickle.HIGHEST_PROTOCOL)