#!/usr/bin/env python3
import sqlalchemy as sa
import sqlalchemy.exc
import sys
from contextlib import contextmanager
from pathlib import Path
import click
import csv

sys.path.append('tm_navigator')
from tm_navigator.models import *

engine = sa.create_engine('postgresql+psycopg2://postgres@localhost/tm_navigator')
sa.orm.configure_mappers()
Session = sa.orm.sessionmaker(bind=engine)

@contextmanager
def session_scope():
    """Provide a transactional scope around a series of operations."""
    session = Session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()

from sqlalchemy_utils import aggregates

aggregates.local_condition = lambda prop, objects: sa.literal(True)

class ListSession(list):
    def execute(self, query):
        return self._session.execute(query)

def update_aggregates(session, *classes_modified):
    ls = ListSession([c() for c in classes_modified])
    ls._session = session
    aggregates.manager.construct_aggregate_queries(ls, None)

def copy_from_csv(session, model, csv_file):
    with session.connection().connection.cursor() as cursor, csv_file.open() as csv_f:
        fieldnames = csv_f.readline()
        cursor.copy_expert('copy {0} ({1}) from stdin with csv'.format(model.__tablename__, fieldnames), csv_f)

    update_aggregates(session, model)

def check_files(directory, expected_names, extension='.csv', cli=True):
    file_names = {p.name for p in directory.iterdir() if p.is_file()}
    expected_files = {t + extension for t in expected_names}

    if expected_files & file_names:
        s = 'Found files {}.'.format(
            ', '.join('"{}"'.format(f) for f in sorted(expected_files & file_names)))
        if cli:
            click.secho(s, fg='green')
        else:
            print(s)
    if expected_files - file_names:
        s = 'Not found files {}.'.format(
            ', '.join('"{}"'.format(f) for f in sorted(expected_files - file_names)))
        if cli:
            click.secho(s, fg='red')
            click.echo('Will try to continue with the files present.')
        else:
            print(s)
            print('Will try to continue with the files present.')

def delete_data_for(session, models, cli=True):
    def delete_data(iter_func=lambda: None):
        for table in reversed(Base.metadata.sorted_tables):
            matching_models = [m for m in models if m.__table__ == table]
            if not matching_models:
                continue
            model = matching_models[0]
            session.query(model).delete()
            iter_func()

    if cli:
        with click.progressbar(label='Deleting data', length=len(models)) as pbar:
            delete_data(lambda: pbar.update(1))
    else:
        delete_data()

def load_data_for(session, models, directory, cli=True):
    def load_data(iter_func=lambda: None):
        for table in Base.metadata.sorted_tables:
            matching_models = [m for m in models if m.__table__ == table]
            if not matching_models:
                continue
            model = matching_models[0]

            file = directory / '{}.csv'.format(model.__tablename__)
            copy_from_csv(session, model, file)
            iter_func()

    if cli:
        with click.progressbar(label='Loading data', length=len(models)) as pbar:
            load_data(lambda: pbar.update(1))

    else:
        load_data()

@click.group()
def cli():
    pass

dir_type = click.Path(exists=True, file_okay=False, resolve_path=True)

@cli.command()
def describe():
    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)
        dses = session.query(DatasetMeta).order_by(DatasetMeta.id).all()
        for ds in dses:
            ds.activate_schemas()
            click.secho('- Dataset #{id}: {title}, {ntm} models'.format(id=ds.id,
                                                                        title=ds.title or 'untitled',
                                                                        ntm=len(ds.topic_models)),
                        fg='blue')
            try:
                click.echo('  Documents: {cnt}'.format(cnt=session.query(Document).count()))
                click.echo('  Terms: ' +
                           ', '.join('{m.count} {m.name} with {nocc} occurrences'.format(m=m, nocc=nocc)
                                     for m, nocc in session.query(Modality, sa.func.count())
                                     .join(Modality.terms).join(Term.documents)
                                     .group_by(Modality).order_by(Modality.name)))
            except sa.exc.ProgrammingError:
                click.echo('  Error - can\'t find data')
                session.rollback()
            click.echo()

            for tm in ds.topic_models:
                tm.activate_schemas()
                click.secho('  - Topic Model #{id}: {title}'.format(id=tm.id, title=tm.title or 'untitled'),
                            fg='blue')
                try:
                    click.echo('    Topics: ' +
                               ', '.join('{cnt} {t} at lvl {lvl}'.format(lvl=lvl, t='background' if bck else '', cnt=cnt)
                                         for lvl, bck, cnt in session.query(Topic.level, Topic.is_background, sa.func.count())
                                         .group_by(Topic.level, Topic.is_background).order_by(Topic.level, Topic.is_background)))
                except sa.exc.ProgrammingError:
                    click.echo('    Error - can\'t find data')
                    session.rollback()
                click.echo()

def add_dataset_():
    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)

        ds = DatasetMeta()
        session.add(ds)
        session.flush()

        ds.activate_schemas()
        Base.metadata.create_all(engine,
                                 tables=map(lambda c: c.__table__, models_dataset))
        return ds.id

@cli.command()
def add_dataset():
    dataset_id = add_dataset_()
    click.echo('Added Dataset #{id}'.format(id=dataset_id))

def load_dataset_(dataset_id, title, directory, cli=False):
    directory = Path(directory)
    target_models = models_dataset

    check_files(directory, [m.__tablename__ for m in target_models], cli=cli)
    if cli:
        click.confirm('Proceeding will overwrite the corresponding data in the database. Continue?',
                      abort=True, default=True)

    models = [m
              for m in target_models
              if (directory / '{}.csv'.format(m.__tablename__)).is_file()]

    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)
        ds = session.query(DatasetMeta).filter_by(id=dataset_id).one()
        ds.activate_schemas()

        if title is not None:
            ds.title = title
        delete_data_for(session, models, cli)
        load_data_for(session, models, directory, cli)

@cli.command()
@click.option('-d', '--dataset-id', type=int, required=True)
@click.option('-t', '--title', type=str)
@click.option('-dir', '--directory', type=dir_type, required=True)
def load_dataset(dataset_id, title, directory):
    load_dataset_(dataset_id, title, directory, cli=True)

def add_topicmodel_(dataset_id):
    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)

        ds = session.query(DatasetMeta).filter_by(id=dataset_id).one()
        tm = TopicModelMeta(dataset=ds)
        session.add(tm)
        session.flush()
        tm.domains.append(TopicModelDomain(domain='{}.'.format(tm.id)))

        tm.activate_schemas()
        Base.metadata.create_all(engine,
                                 tables=map(lambda c: c.__table__, models_topic + models_assessment))

        return tm.id

@cli.command()
@click.option('-d', '--dataset-id', type=int, required=True)
def add_topicmodel(dataset_id):
    topicmodel_id = add_topicmodel_(dataset_id)
    click.echo('Added Topic Model #{id} for Dataset #{did}'.format(id=topicmodel_id, did=dataset_id))

def load_topicmodel_(topicmodel_id, title, directory, cli=False):
    directory = Path(directory)
    target_models = models_topic

    check_files(directory, [m.__tablename__ for m in target_models], cli=cli)
    if cli:
        click.confirm('Proceeding will overwrite the corresponding data in the database. Continue?',
                      abort=True, default=True)

    models = [m
              for m in target_models
              if (directory / '{}.csv'.format(m.__tablename__)).is_file()]

    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)
        tm = session.query(TopicModelMeta).filter_by(id=topicmodel_id).one()
        tm.activate_schemas()

        if title is not None:
            tm.title = title
        delete_data_for(session, models, cli)
        load_data_for(session, models, directory, cli)

@cli.command()
@click.option('-m', '--topicmodel-id', type=int, required=True)
@click.option('-t', '--title', type=str)
@click.option('-dir', '--directory', type=dir_type, required=True)
def load_topicmodel(topicmodel_id, title, directory):
    load_topicmodel_(topicmodel_id, title, directory, cli=True)

def dump_assessments_(topicmodel_id, directory):
    directory = Path(directory)
    with session_scope() as session:
        SchemaMixin.activate_public_schema(session)
        tm = session.query(TopicModelMeta).filter_by(id=topicmodel_id).one()
        tm.activate_schemas()

        assessments = session.query(ATopic).all()
        topic_count = session.query(Topic).count() - 1

        grades = {}
        for assessment in assessments:
            grades[assessment.topic_id] = assessment.value
        with open(directory.joinpath('topic_assessments.csv'), 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['topic_id', 'value'])
            for idx, grade in grades.items():
                writer.writerow([idx, grade])

@cli.command()
@click.option('-m', '--topicmodel-id', type=int, required=True)
@click.option('-dir', '--directory', type=dir_type, required=True)
def dump_assessments(topicmodel_id, directory):
    dump_assessments_(topicmodel_id, directory)

if __name__ == '__main__':
    cli()


import unittest
from hamcrest import assert_that, instance_of, equal_to, raises, calling
from app.exiftool.s3.object import Object
from app.exiftool.s3.object_iterator import ObjectIterator
from tests.config import Config as TestConfig

class TestS3ObjectIterator(unittest.TestCase):
    def setUp(self) -> None:
        self.config = TestConfig()
        self.config.patcher.environment()
        self.addCleanup(self.config.patcher.stop())

    def test_next_success(self) -> None:
        with self.config.use_cassette('s3_object_iterator_succeeded_pdf.yml'):
            object_iterator = ObjectIterator(
                Object('test-bucket', 'test/test.pdf')
            )

            assert_that(
                object_iterator.__next__(),
                instance_of(bytes)
            )
            assert_that(object_iterator.first_byte_pos, equal_to(0))
            assert_that(object_iterator.last_byte_pos, equal_to(5300))

    def test_next_failure(self) -> None:
        with self.config.use_cassette('s3_object_iterator_succeeded_pdf.yml'):
            object_iterator = ObjectIterator(
                Object('test-bucket', 'test/test.pdf')
            )
            object_iterator.__next__()

            assert_that(
                calling(object_iterator.__next__),
                raises(StopIteration)
            )  # type: ignore


from pyspark.context import SparkContext
import itertools
import collections
from collections import Counter
import time

from operator import add
import os
import sys

def gencomb(freqitem,lenfreqitem):
    cand= []
    for x in range(0,lenfreqitem-1):
         first = freqitem[x]
         firstelem = sorted(first)
         for y in range(x+1,lenfreqitem):
              second = freqitem[y]
              secelem= sorted(second)
              newItem = []
              if firstelem[0:len(firstelem)-1] ==  secelem[0:len(secelem)-1] :
                   newItem=list(firstelem)
                   newItem.append(secelem[len(secelem)-1])
                   newItem.sort()
                   cand.append((newItem))
    #print("candidate ",cand)
    return cand

def checksub(checksub, prev, len, initbasket, ktuple):
    store = Counter()
    candidates = []
    for k in checksub:
        item = set(k)
        count = 0
        for x in prev:
            if set(x).issubset(item):
                count += 1
            if count == ktuple:
                candidates.append(tuple(k))
                break

    candidates = list(dict.fromkeys(candidates))

    for basket in initbasket:
        bas = set(basket)
        for candidate in candidates:
            temp = set(candidate)
            if temp.issubset(bas):
                store[tuple(candidate)] += 1
    return store

def prune(checksubs, threshold):
    freq = []
    for x in checksubs:
        if checksubs[x] >= threshold:
            freq.append(tuple(x))
    return sorted(freq)

def apriori(basket, basket_count, sup):
    initbasket = list(basket)
    support = float(sup)
    partition_baskets = len(initbasket)
    #print("basket ", partition_baskets)
    total_baskets = basket_count
    threshold = support * (float(partition_baskets) / float(total_baskets))
    #print("THRESHOLD ", threshold)
    final_itemsets = []
    single_count = {}
    freqsingle = []
    single_print = []
    # SINGLTON

    for bas in initbasket:
        for item in bas:
            if item in single_count:
                single_count[item] += 1
            else:
                single_count[item] = 1

    for x in single_count:
        if single_count[x] >= threshold:
            freqsingle.append(x)
            single_print.append((x,))
    single = sorted(tuple(freqsingle))
    single_print = sorted(single_print)
    final_itemsets = single
    #print("single  ",len(single))
# PAIRS

    res = list(itertools.combinations(single, 2))
    #print("res ",len(res))
    paircount = collections.defaultdict(int)
    freqpair = []
    cand = []

    for basket in initbasket:
        bas = set(basket)
        for candidate in res:
            temp= set(candidate)
            if temp.issubset(bas):
                if paircount.get(candidate):
                    paircount[candidate] += 1
                else:
                    paircount[candidate] = 1

    for item,count in paircount.items():
        if count >= threshold:
            freqpair.append(tuple(sorted(item)))

    pair = sorted(freqpair)
    pair = list(dict.fromkeys(pair))
    #print("pair   ",len(pair))
    final_itemsets = final_itemsets + pair
    single_print = single_print + pair
    # print("pair   ",single_print)
    frequent = pair
    k = 3
    while frequent != []:

        gencombination = gencomb(frequent, len(frequent))
        # print(k," itemsets",len(gencombination))
        checksubsets = checksub(gencombination, frequent, len(gencombination), initbasket, k)
        # print("candidate ", len(checksubsets))
        freqitemsets = prune(checksubsets, threshold)
        #print(k, " -tuple freqitem ", len(freqitemsets))
        frequent = list(freqitemsets)
        if frequent != []:
            final_itemsets = final_itemsets + frequent
            single_print = single_print + frequent
        k += 1
    # print("Final itemsets ", single_print)

    yield (single_print)

def countmap2(bt):
    initbasket=list(bt)
    count = Counter()

    for basket in initbasket:
        bas=set(basket)
        for item in val:
            if item in bas:
                count[item] += 1
            else:
                temp = set(item)
                if temp.issubset(bas):
                    count[tuple(item)] += 1
    #print("countint ",list(count.items()))

    yield (list(count.items()))

if __name__ == "__main__":

    filterthreshold = sys.argv[1]
    support = sys.argv[2]
    input_path = sys.argv[3]
    output_path = sys.argv[4]

    sc= SparkContext("local[*]")
    sc.setLogLevel("ERROR")
    start = time.time()
    lines = sc.textFile(input_path)
    header = lines.first()
    RDD= lines.filter(lambda x: x != header).map(lambda x: tuple(x.split(",")))
    basket1= RDD.map(lambda x: (x[0],[x[1]])).reduceByKey(lambda x,y: x + y)
    bas=basket1.filter(lambda x: len(x[1]) > int(filterthreshold)).map(lambda x: (x[0],x[1]))
    bt=bas.values()
    basket_count = bt.count()

    mapbasket1 = bt.mapPartitions(lambda x: apriori(x,basket_count,support)).flatMap(lambda x: x)
    phase1_map= mapbasket1.map(lambda itemsets: (itemsets,1))

    phase1_reduce= phase1_map.reduceByKey(lambda x,y : 1).sortBy(keyfunc= lambda x: (len(x[0]),x[0]))
    val= phase1_reduce.keys().collect()
    #print("phase1reduce ", len(val))

    with open(output_path, 'w') as fileop:
        fileop.write("Candidates: \n")

        old = 1
        storestr = ""
        for index in range(len(val)):
            element = val[index]

            if index == len(val) - 1:
                fileop.write(storestr.replace(")(", "),("))
                if len(element) > old:
                    fileop.write("\n\n")
                    fileop.write(str(element))
                    break
                else:
                    fileop.write(","+str(element))
                    break

            if len(element) > old:
                fileop.write(storestr.replace(")(", "),("))
                fileop.write("\n\n")
                old += 1
                storestr = ""
                fileop.write(str(element))
            else:
                if len(element) == 1:
                    storestr += str(element).replace(",)", ")")
                else:
                    storestr += "," + str(element)
        fileop.close()

    phase2_map = bt.mapPartitions(countmap2).flatMap(lambda x: x).reduceByKey(add).sortBy(keyfunc= lambda x: (len(x[0]),x[0]))

    phase2_reduce = phase2_map.filter(lambda x: x[1] >= float(support)).map(lambda x: x[0])
    res = phase2_reduce.collect()
    #print("res ", len(res))

    with open(output_path, 'a') as fileop:
        fileop.write("\n\nFrequent Itemsets: \n")

        old = 1
        storestr = ""
        for index in range(len(res)):
            element = res[index]

            if index == len(res) - 1:
                fileop.write(storestr.replace(")(", "),("))
                if len(element) > old:
                    fileop.write("\n\n")
                    fileop.write(str(element))
                    break
                else:
                    fileop.write(","+str(element))
                    break

            if len(element) > old:
                fileop.write(storestr.replace(")(", "),("))
                fileop.write("\n\n")
                old += 1
                storestr = ""
                fileop.write(str(element))
            else:
                if len(element) == 1:
                    storestr += str(element).replace(",)", ")")
                else:
                    storestr += "," + str(element)
        fileop.close()
    end = time.time()
    print("Duration:", end - start)


import random
import socket

buf = 1024
timeout = 40
class Packet:

    def __init__(self, seq_n, is_ack, data, checksum=None):
        self.seq_n = seq_n
        self.is_ack = is_ack
        self.data = data
        self.checksum = checksum

        if is_ack == "True":
            is_ack = True
        if is_ack == "False":
            is_ack = False

        if checksum is None:
            self.checksum = self.real_checksum()

    def reading_size(self):
        _checksum = self.real_checksum()
        packet_return = (str(self.seq_n) + "|" + str(_checksum) + "|" + str(self.is_ack) + "|")
        return len(packet_return.encode('utf-8')) + 16

    def make_packet(self):
        _checksum = self.real_checksum()
        packet_return = (str(self.seq_n) + "|" + str(_checksum) + "|" + str(self.is_ack) + "|" + str(self.data))
        return packet_return

    def real_checksum(self):
        data = str(self.seq_n) + str(self.is_ack) + str(self.data)
        data = data.encode()

        polynomial = 0x1021
        crc = 0xFFFF
        for byte in data:
            crc ^= (byte << 8)
            for _ in range(8):
                if (crc & 0x8000):
                    crc = (crc << 1) ^ polynomial
                else:
                    crc = (crc << 1)
        return crc & 0xFFFF

    def is_corrupt(self):
        return self.real_checksum() != self.checksum

def extract_packet(string_packet):
    seq_num, checksum_, is_ack, data = string_packet.decode().split("|", 3)
    return Packet(int(seq_num), is_ack, data, checksum=int(checksum_))

def send_packet(sock, packet, addr):
    #print(f"Enviando: {packet.seq_n}")
    #print(" " + str(len(packet.make_packet().encode('utf-8'))) + " bytes")
    sock.sendto(packet.make_packet().encode(), addr)

def send_ack(sock, seq_num, addr):
    packet = Packet(seq_num, True, 0, 0)
    #print(f"Enviando ACK: {seq_num}")
    sock.sendto(packet.make_packet().encode(), addr)

def wait_for_ack(sock, expected_ack):
    try:
        ack = Packet(0, True, "")
        data, _ = sock.recvfrom(buf - ack.reading_size())
        ack = extract_packet(data)
        if ack.is_ack and expected_ack == ack.seq_n:
            #print(f"ACK recebido: {ack.seq_n}")
            return True
        else:
            #print(f"[ERRO] ACK: {ack.seq_n}, esperado: {expected_ack}")
            return False
    except socket.timeout:
        #print(f"Tempo limite excedido: {timeout} segundos.")
        return False

def packet_loss(probability):
    return random.random() < probability


import asyncio

from django.core.cache import cache

from apps.integration.models import KotakNeoApi as KotakNeoApiModel
from apps.integration.utils import divide_and_list, get_option_ltp
from apps.integration.utils.broker.dummy import DummyApi
from apps.integration.utils.broker.kotak_neo import KotakNeoApi
from apps.trade.models import DeployedOptionStrategy as DeployedOptionStrategyModel

class StrategyOrder(object):
    def __init__(self, opt_strategy: DeployedOptionStrategyModel):
        self.opt_strategy = opt_strategy
        self.symbol = opt_strategy.instrument.symbol
        self.websocket_ids = self.opt_strategy.websocket_ids.split(",")
        self.strategy_id = self.opt_strategy.pk
        self.str_strategy_id = str(self.opt_strategy.pk)
        self.options = self.opt_strategy.options.copy()
        self.parameters = [
            p.parameters
            for p in self.opt_strategy.parameters.filter(is_active=True).order_by(
                "name"
            )
        ]
        self.parameters_len = len(self.parameters) or 1
        self.entry_type = self.opt_strategy.strategy_type
        self.exit_type = "BUY" if self.entry_type == "SELL" else "SELL"
        self.strategy_tradingsymbol_cache = f"TRADINGSYMBOLS_{self.opt_strategy.pk}"
        self.broker = self.opt_strategy.broker
        self.slippage = float(self.opt_strategy.slippage)
        user_params = (
            cache.get("DEPLOYED_STRATEGIES", {})
            .get(self.str_strategy_id, {})
            .get("user_params", [])
        )
        if user_params and cache.get(self.strategy_tradingsymbol_cache):
            self.user_params = user_params
            self.entered = True
        else:
            self.entered = False
            self.user_params = []
            for user in self.opt_strategy.users.filter(is_active=True).order_by(
                "order_seq"
            ):
                self.user_params.append(
                    {
                        "broker_api": user.broker_api,
                        "user": user.broker_api.user,
                        "broker": user.broker_api.broker,
                        "quantity_multiple": [
                            item * self.opt_strategy.lot_size
                            for item in divide_and_list(self.parameters_len, user.lots)
                        ],
                    }
                )

    async def get_kotak_neo_parameters(self, user):
        row = KotakNeoApiModel.objects.get(broker_api__user=user)

        return {
            "username": user.username,
            "headers": {
                "Authorization": f"Bearer {row.access_token}",
                "Sid": row.sid,
                "Auth": row.auth,
                "neo-fin-key": row.neo_fin_key,
                "Content-Type": "application/x-www-form-urlencoded",
            },
            "query_params": {"sId": row.hs_server_id},
        }

    async def quantity_wise_strategy_map(self, user_params):
        quantity_map_list = []

        for idx in range(self.parameters_len):
            quantity_map_list.append(
                {
                    user_param["user"].username: user_param["quantity_multiple"][idx]
                    for user_param in user_params
                }
            )

        return quantity_map_list

    async def place_dummy_order(self, buy_pending, sell_pending, user_params):
        quantity_map_list = await self.quantity_wise_strategy_map(user_params)
        users = [user_param["user"] for user_param in user_params]

        if not user_params:
            return

        dummyapi = DummyApi(users=users)

        if buy_pending:
            await asyncio.gather(
                *[
                    dummyapi.option_place_and_chase_order(
                        symbol=self.symbol,
                        tradingsymbol=row["tradingsymbol"],
                        quantity_map=quantity_map_list[row.get("idx", 0)],
                        transaction_type="BUY",
                        price=get_option_ltp(
                            symbol=self.symbol,
                            tradingsymbol=row["tradingsymbol"],
                            websocket_id=row["websocket_id"],
                        ),
                        tag_prefix=self.str_strategy_id,
                        websocket_id=row["websocket_id"],
                    )
                    for row in buy_pending
                ]
            )

        if sell_pending:
            await asyncio.gather(
                *[
                    dummyapi.option_place_and_chase_order(
                        symbol=self.symbol,
                        tradingsymbol=row["tradingsymbol"],
                        quantity_map=quantity_map_list[row.get("idx", 0)],
                        transaction_type="SELL",
                        price=get_option_ltp(
                            symbol=self.symbol,
                            tradingsymbol=row["tradingsymbol"],
                            websocket_id=row["websocket_id"],
                        ),
                        tag_prefix=self.str_strategy_id,
                        websocket_id=row["websocket_id"],
                    )
                    for row in sell_pending
                ]
            )

    async def place_kotak_neo_order(self, buy_pending, sell_pending, user_params):
        quantity_map_list = await self.quantity_wise_strategy_map(user_params)
        users = [
            await self.get_kotak_neo_parameters(user_param["user"])
            for user_param in user_params
        ]
        if not users:
            return

        knapi = KotakNeoApi(users=users)

        if buy_pending:
            await asyncio.gather(
                *[
                    knapi.option_place_and_chase_order(
                        symbol=self.symbol,
                        tradingsymbol=row["tradingsymbol"],
                        quantity_map=quantity_map_list[row.get("idx", 0)],
                        transaction_type="BUY",
                        expected_price=get_option_ltp(
                            symbol=self.symbol,
                            tradingsymbol=row["tradingsymbol"],
                            websocket_id=row["websocket_id"],
                        )
                        + self.slippage,
                        tag_prefix=self.str_strategy_id,
                        websocket_id=row["websocket_id"],
                    )
                    for row in buy_pending
                ]
            )

        if sell_pending:
            await asyncio.gather(
                *[
                    knapi.option_place_and_chase_order(
                        symbol=self.symbol,
                        tradingsymbol=row["tradingsymbol"],
                        quantity_map=quantity_map_list[row.get("idx", 0)],
                        transaction_type="SELL",
                        expected_price=get_option_ltp(
                            symbol=self.symbol,
                            tradingsymbol=row["tradingsymbol"],
                            websocket_id=row["websocket_id"],
                        )
                        - self.slippage,
                        tag_prefix=self.str_strategy_id,
                        websocket_id=row["websocket_id"],
                    )
                    for row in sell_pending
                ]
            )

    async def place_order(
        self,
        sell_pending: list,
        buy_pending: list,
        user_params: list | None = None,
    ):
        if buy_pending or sell_pending:
            if not user_params:
                user_params = self.user_params

            if self.broker == "kotak_neo":
                kotak_neo_users_params = [
                    user_param
                    for user_param in user_params
                    if user_param["broker"] == "kotak_neo"
                ]

                dummy_users_params = [
                    user_param
                    for user_param in user_params
                    if user_param["broker"] == "dummy"
                ]

                await asyncio.gather(
                    self.place_dummy_order(
                        buy_pending,
                        sell_pending,
                        user_params=dummy_users_params,
                    ),
                    self.place_kotak_neo_order(
                        buy_pending,
                        sell_pending,
                        user_params=kotak_neo_users_params,
                    ),
                )


from threading import Thread, Lock
from math import ceil

lst = []
NO_THREADS = 8
total_sum = 0
lock = Lock()

class MyThread(Thread):
    def __init__(self, index):
        Thread.__init__(self)
        self.index = index

    def run(self):
        global total_sum
        start = self.index * ceil(len(lst) / NO_THREADS)
        end = min(len(lst), ((self.index + 1) * ceil(len(lst) / NO_THREADS)))

        for i in range(start, end):
            lock.acquire()
            total_sum += lst[i]
            lock.release()
            '''
            # alternativ
            with lock:
                total_sum += lst[i]
            '''

def main():
    global lst

    lst = [x for x in range(1, 101)]

    threads = []
    for i in range(0, NO_THREADS):
        threads.append(MyThread(i))
        threads[-1].start()

    for t in threads:
        t.join()
    print(total_sum)

if __name__ == '__main__':
    main()


import yaml
import os
import logging
from weight import Weight
from schema import Schema, SchemaError, Optional
from typing import Union

_LOGGER = logging.getLogger(__name__)
_LOGGER.setLevel(logging.DEBUG)

class KB_Chaos:
    def __init__(self, chaos_path):
        self.chaos_path = chaos_path
        self.last_chaos = None

    def is_instance_related(self):
        """Check whether chaos is instance related"""

        if len(os.listdir(self.chaos_path)) == 0:
            _LOGGER.error("No chaos found in {}".format(self.chaos_path))
            return True, None

        for chaos in os.listdir(self.chaos_path):
            if chaos.endswith(".yaml"):
                f = open(self.chaos_path + "/" + chaos)
                data = f.read()
                f.close()
                data = yaml.safe_load(data)
                if self.last_chaos is None:
                    self.last_chaos = data

                else:
                    if data["anomalies"] == self.last_chaos["anomalies"]:
                        continue
                    else:
                        return True, data

        if len(os.listdir(self.chaos_path)) <= 2:
            return True, data

        return False, data

class KB:
    def __init__(self) -> None:

        self.kb = None
        self.hierarchy = {0: "chaos_type", 1: "chaos"}
        self.metrics = []
        self.traces = []
        self.logs = []
        self.cmds = []

        self.type_metrics = []
        self.type_traces = []
        self.type_logs = []
        self.type_cmds = []

        self.metrics_score = None
        self.traces_score = None
        self.logs_score = None
        self.cmds_score = None

        self.type_metrics_score = None
        self.type_traces_score = None
        self.type_logs_score = None
        self.type_cmds_score = None

    def load(self, kb_path: str) -> Union[dict, None]:
        """Load knowledge base

        Args:
            kb_path (str): Knowledge base path

        Raises:
            Exception: Knowledge base check

        Returns:
            dict: Knowledge base
        """

        if type(kb_path) is str:
            f = open(kb_path)
            data = f.read()
            f.close()
            self.kb = yaml.safe_load(data)
        elif type(kb_path) is dict:
            self.kb = kb_path

        is_checked = self.check_kb()

        if is_checked:
            self.score_fingerprint()
            return self.kb
        else:
            raise Exception("Knowledge Base check failed")

    def check_kb(self) -> bool:
        """Check knowledge base config

        Raises:
            se: Schema error

        Returns:
            bool: check result
        """
        if self.kb is None:
            _LOGGER.error("Knowledge Base is not loaded")
            return False

        anomaly_schema = [{"index": int, "action": str, Optional("order"): int}]
        custom_metrics_schema = {
            Optional("network"): anomaly_schema,
            Optional("cpu"): anomaly_schema,
            Optional("memory"): anomaly_schema,
            Optional("io"): anomaly_schema,
            Optional("container"): anomaly_schema,
            Optional("mongo"): anomaly_schema,
            Optional("mysql"): anomaly_schema,
            Optional("icmp"): anomaly_schema,
            Optional("time"): anomaly_schema,
            Optional("jvm"): anomaly_schema,
            Optional("http"): anomaly_schema,
        }

        custom_traces_schema = {
            Optional("onehop"): anomaly_schema,
        }

        custom_logs_schema = {
            Optional("pod"): anomaly_schema,
        }

        custom_cmds_schema = {
            Optional("config"): anomaly_schema,
            Optional("exec"): anomaly_schema,
        }

        custom_schema = [
            {
                "index": int,
                "experiment": str,
                "instance_related": bool,
                Optional("order"): bool,
                "anomalies": {
                    Optional("metrics"): custom_metrics_schema,
                    Optional("traces"): custom_traces_schema,
                    Optional("logs"): custom_logs_schema,
                    Optional("cmds"): custom_cmds_schema,
                },
            }
        ]

        config_schema = Schema(
            {
                Optional("network"): custom_schema,
                Optional("pod"): custom_schema,
                Optional("stress"): custom_schema,
                Optional("time"): custom_schema,
                Optional("jvm"): custom_schema,
                Optional("dns"): custom_schema,
                Optional("http"): custom_schema,
                Optional("io"): custom_schema,
                Optional("config"): custom_schema,
            }
        )

        try:
            config_schema.validate(self.kb)
            _LOGGER.info("Configuration is valid.")
        except SchemaError as se:
            raise se

        return True

    def score_fingerprint(self):
        """Score fingerprint"""

        # Two hierarchies for our experiment

        chaos_types = self.kb.keys()

        for chaos_type in chaos_types:

            type_metrics = []
            type_traces = []
            type_logs = []
            type_cmds = []
            for chaos in self.kb[chaos_type]:
                anomalies = chaos["anomalies"]
                metrics = (
                    anomalies["metrics"] if "metrics" in anomalies else None
                )
                traces = anomalies["traces"] if "traces" in anomalies else None
                logs = anomalies["logs"] if "logs" in anomalies else None
                cmds = anomalies["cmds"] if "cmds" in anomalies else None

                (
                    metrics_instance,
                    traces_instance,
                    logs_instance,
                    cmds_instance,
                ) = self.analyse(metrics, traces, logs, cmds)

                type_metrics += metrics_instance
                type_traces += traces_instance
                type_logs += logs_instance
                type_cmds += cmds_instance

            self.type_metrics.append(type_metrics) if type_metrics else None
            self.type_traces.append(type_traces) if type_traces else None
            self.type_logs.append(type_logs) if type_logs else None
            self.type_cmds.append(type_cmds) if type_cmds else None

        for (data, score) in zip(
            [
                self.metrics,
                self.traces,
                self.logs,
                self.cmds,
                self.type_metrics,
                self.type_traces,
                self.type_logs,
                self.type_cmds,
            ],
            [
                "metrics_score",
                "traces_score",
                "logs_score",
                "cmds_score",
                "type_metrics_score",
                "type_traces_score",
                "type_logs_score",
                "type_cmds_score",
            ],
        ):
            weight = Weight(data)
            weighted_score = weight()
            max_score = max(weighted_score.values())
            for key in weighted_score:
                weighted_score[key] = weighted_score[key] / max_score
            setattr(self, score, weighted_score)

    def analyse(
        self, metrics: list, traces: list, logs: list, cmds: list
    ) -> tuple:
        """Analyse metrics, traces, logs, cmds

        Args:
            metrics (list): metrics
            traces (list): traces
            logs (list): logs
            cmds (list): commands

        Returns:
            tuple: metrics, traces, logs, cmds
        """

        metrics_instance = self.analyse_fingerprint(metrics, "metrics")
        traces_instance = self.analyse_fingerprint(traces, "traces")
        logs_instance = self.analyse_fingerprint(logs, "logs")
        cmds_instance = self.analyse_fingerprint(cmds, "cmds")

        return metrics_instance, traces_instance, logs_instance, cmds_instance

    def analyse_fingerprint(
        self, fingerprint: list, target_type: str = ""
    ) -> list:
        """Analyse fingerprint individually

        Args:
            fingerprint (list): Fingerprint
            target_type (str, optional): Fingerprint type. Defaults to "".

        Returns:
            list: Rename instances
        """

        if fingerprint is None or target_type == "":
            _LOGGER.info("No {} found in Knowledge Base".format(target_type))
            return []

        types = fingerprint.keys()
        new_instance = []
        for one_type in types:
            for clue in fingerprint[one_type]:
                idx = clue["index"]
                action = clue["action"]
                clue_name = one_type + "-" + str(idx) + "-" + action
                new_instance.append(clue_name)

        if new_instance:
            if target_type == "metrics":
                self.metrics.append(new_instance)
            elif target_type == "traces":
                self.traces.append(new_instance)
            elif target_type == "logs":
                self.logs.append(new_instance)
            elif target_type == "cmds":
                self.cmds.append(new_instance)
            self.logs.append(new_instance)

        return new_instance


from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.utils.dummy_pt_objects import PreTrainedModel
from transformers.data.processors.utils import InputExample, InputFeatures

from .prompts import TemplateGenerator, VerbalizerGenerator
from openprompt import PromptDataLoader, PromptForClassification
from openprompt.prompts import ManualVerbalizer, LMBFFTemplate, ManualTemplate
from typing import List, Optional, Dict, Union
from . import Template, Verbalizer, PromptDataLoader
import copy
import warnings
from .trainer import ClassificationRunner
from yacs.config import CfgNode
from openprompt.utils.logging import logger
from openprompt.utils.cuda import model_to_device

class LMBFFClassificationRunner:
    def __init__(self,
                train_dataset: List[InputExample],
                valid_dataset: List[InputExample],
                test_dataset: List[InputExample],
                model: PreTrainedModel,
                tokenizer: PreTrainedTokenizer,
                template_generator_tokenizer: PreTrainedTokenizer,
                initial_template: Union[LMBFFTemplate, ManualTemplate],
                initial_verbalizer: ManualVerbalizer = None,
                template_generator: TemplateGenerator = None,
                verbalizer_generator: VerbalizerGenerator = None,
                config: CfgNode = None):
        r"""
        This class implements the LM-BFF in paper (https://arxiv.org/pdf/2012.15723.pdf)
        """
        self.model = model
        self.tokenizer = tokenizer
        self.template_generator_tokenizer = template_generator_tokenizer
        self.max_epoch = config.classification.generation_epoch
        self.train_dataset = train_dataset
        self.valid_dataset = valid_dataset
        self.test_dataset = test_dataset
        self.template_generator = template_generator
        self.verbalizer_generator = verbalizer_generator
        self.config = config
        self.metric = self.config.classification.metric[0]
        self.initial_template = initial_template
        self.initial_verbalizer = initial_verbalizer
        self.auto_t = config.classification.auto_t
        self.auto_v = config.classification.auto_v

        if self.auto_t:
            if not (initial_template and initial_verbalizer):
                raise ValueError("initial template and verbalizer must not be None when both auto_t and auto_v are performed")
            if not isinstance(initial_template, LMBFFTemplate):
                raise ValueError("To perform template search, initial_template must be LMBFFTemplate, however {} is provided".format(type(initial_template.__class__.__name__)))
        elif self.auto_v:
            if initial_verbalizer is not None:
                warnings.warn("only auto_v is True, the given initial_verbalizer is ignored")
            if isinstance(initial_template, LMBFFTemplate):
                raise ValueError("To perform template search, initial_template must be ManualTemplate, however LMBFFTemplate is provided")
        else:
            warnings.warn("auto_t and auto_v are both False, the trainer will degenerate to a simple classification trainer")

    def _auto_t(self, dataset, template, verbalizer):
        logger.info("performing auto-t...")
        if self.template_generator is None:
            raise ValueError("no template_generator available")
        dataloader = PromptDataLoader(dataset, template, self.template_generator_tokenizer, batch_size=len(dataset)) # register all data at once
        for data in dataloader:
            data = data.to("cuda:{}".format(self.config.environment.local_rank))
            self.template_generator.register_buffer(data.input_ids, data.attention_mask, data.label)
        template_texts = self.template_generator.generate() # [['text_a', '<mask>', ...]]
        best_template_text = self._get_best_template_text(template_texts, verbalizer)
        return best_template_text

    def _auto_v(self, dataset, template, verbalizer, batch_size=32):
        logger.info("performing auto-v...")
        if self.verbalizer_generator is None:
            raise ValueError("no verbalizer_generator available")
        dataloader = PromptDataLoader(dataset, template, self.tokenizer, batch_size=batch_size)
        for data in dataloader:
            data = template.process_batch(data)
            data = data.to("cuda:{}".format(self.config.environment.local_rank))
            self.verbalizer_generator.register_buffer(data)
        label_words_list = self.verbalizer_generator.generate() # List[List[str]]
        best_label_words = self._get_best_label_words(label_words_list, template, verbalizer)
        return best_label_words

    def run(self):
        '''
        if both auto_v and auto_t are set to True, perform auto_t first and then auto_v
        '''

        best_template = self.initial_template
        best_verbalizer = self.initial_verbalizer
        if self.auto_t:
            best_template_text = self._auto_t(self.train_dataset, self.initial_template, best_verbalizer)
            best_template = ManualTemplate(self.tokenizer, best_template_text)
        if self.auto_v:
            best_label_words = self._auto_v(self.train_dataset, best_template, best_verbalizer)
            best_verbalizer.label_words = best_label_words

        train_dataloader = PromptDataLoader(self.train_dataset, best_template, self.tokenizer)
        valid_dataloader = PromptDataLoader(self.valid_dataset, best_template, self.tokenizer)
        test_dataloader = PromptDataLoader(self.test_dataset, best_template, self.tokenizer)
        model = PromptForClassification(copy.deepcopy(self.model), best_template, best_verbalizer)
        model = model_to_device(model, self.config.environment)
        runner = ClassificationRunner(model, train_dataloader, valid_dataloader, test_dataloader, config=self.config)
        runner.run()

    def _get_best_template_text(self, template_texts_candidates, verbalizer):
        best_metrics = 0.0
        best_template_text = None
        for template_text in template_texts_candidates:
            print(template_text)
            template = ManualTemplate(self.tokenizer, template_text)
            train_dataloader = PromptDataLoader(self.train_dataset, template, self.tokenizer)
            valid_dataloader = PromptDataLoader(self.valid_dataset, template, self.tokenizer)
            score = self._train_eval(template, verbalizer, train_dataloader, valid_dataloader)
            if score > best_metrics:
                best_template_text = template_text
        return best_template_text

    def _get_best_label_words(self, verbalizer_labelwords_candidates, template, verbalizer):
        current_verbalizer = copy.deepcopy(verbalizer)
        best_metrics = 0.0
        best_label_words = None
        for label_words in verbalizer_labelwords_candidates:
            current_verbalizer.label_words = label_words
            train_dataloader = PromptDataLoader(self.train_dataset, template, self.tokenizer)
            valid_dataloader = PromptDataLoader(self.valid_dataset, template, self.tokenizer)
            score = self._train_eval(template, current_verbalizer, train_dataloader, valid_dataloader)
            if score > best_metrics:
                best_label_words = label_words
        return best_label_words

    def _train_eval(self, template, verbalizer, train_dataloader, valid_dataloader):
        model = PromptForClassification(copy.deepcopy(self.model), template, verbalizer)
        model = model_to_device(model, self.config.environment)
        runner = ClassificationRunner(model, train_dataloader, valid_dataloader, config=self.config)
        best_score = 0.0
        for epoch in range(self.max_epoch):
            runner.train_epoch(epoch)
            scores = runner.evaluate(valid_dataloader, 'Valid')
            score = scores[self.metric]
            if score > best_score:
                best_score = score
        return best_score


# Faça um programa para o cálculo de uma folha de pagamento,
# sabendo que os descontos são do Imposto de Renda, que depende
# do salário bruto (conforme tabela abaixo) e 3% para o Sindicato
# e que o FGTS corresponde a 11% do Salário Bruto, mas não é descontado
# (é a empresa que deposita). O Salário Líquido corresponde ao Salário Bruto menos os descontos.
# O programa deverá pedir ao usuário o valor da sua hora e a quantidade de horas trabalhadas no mês.
#       Desconto do IR:
#   *   Salário Bruto até 900 (inclusive) - isento
#   *   Salário Bruto até 1500 (inclusive) - desconto de 5%
#   *   Salário Bruto até 2500 (inclusive) - desconto de 10%
#   *   Salário Bruto acima de 2500 - desconto de 20% Imprima na tela as informações,
# dispostas conforme o exemplo abaixo. No exemplo o valor da hora é 5 e a quantidade de hora é 220.
print('-' * 42)
valorHora = float(input('Qual o valor da sua hora de trabalho? R$'))
print('-' * 42)
horasTrabalhadas = int(input('Quantas horas você trabalhou este mês? '))
print('-' * 42)
salárioBruto = valorHora * horasTrabalhadas
print(f'{"Salário Bruto":30}: R${salárioBruto:>8.2f}')

if 900 < salárioBruto <= 1500:
    iRPercentual = 0.05
elif 1500 < salárioBruto <=2500:
    iRPercentual = 0.1
elif salárioBruto > 2500:
    iRPercentual = 0.2
else:
    iRPercentual = 0
iR = salárioBruto * iRPercentual

if iRPercentual != 0:
    print(f'{f"(-) IR ({int(iRPercentual * 100)}%)":30}: R${iR:>8.2f}')
else:
    print(f'{"(-) IR (Isento)":30}: R${iR:8.2f}')

sindicato = salárioBruto * 0.03
print(f'{"(-) Sindicato (3%)":30}: R${sindicato:>8.2f}')
iNSS = salárioBruto * 0.1
print(f'{"(-) INSS (10%)":30}: R${iNSS:>8.2f}')
fGTS = salárioBruto * 0.11
print(f'{"FGTS (11%)":30}: R${fGTS:>8.2f}')
totalDosDescontos = iR + iNSS +sindicato
print(f'{"Total de descontos":30}: R${totalDosDescontos:>8.2f}')
salárioLiquido = salárioBruto - totalDosDescontos
print(f'{"Salário Liquido":30}: R${salárioLiquido:>8.2f}')
print('-' * 42)


import os
import random

from urllib.parse import quote
from typing import List, Optional
from dataclasses import dataclass, field
from enum import Enum

from jinja2 import Template

SEAT_NUMBER_MAX_LENGTH = 3

# ascii characters are prohibited :D
ru_alphabet_lower = 'абвгдеёжзиклмопрстуфхцчшщэюя'
names = ['Сергій', 'Іван', 'Микола', 'Софія', 'Дмитро', 'Ганна', 'Надія', 'Віктор', 'Василь']
surnames = ['Іванов', 'Морозов', 'Третяк', 'Проценко', 'Сорокін', 'Дубілет', 'Зеленський']

class WagonType(Enum):
    K = 'К'  # Купе
    P = 'П'  # Плацкарт
    L = 'Л'  # Люкс

    def __str__(self):
        return self.value

@dataclass
class Seat:
    number: str
    first_name: str
    last_name: str

    def __str__(self) -> str:
        return f'Seat <number: {self.number}>'

    def __post_init__(self):
        if 3 < len(self.number) < 3:
            self._normalize_seat_number()

    def _normalize_seat_number(self):
        self.number = self.number.zfill(SEAT_NUMBER_MAX_LENGTH)

@dataclass
class TrainOrder:
    from_station_id: int
    to_station_id: int
    train_code: str
    date: str
    wagon_number: int
    wagon_class: str
    wagon_type: WagonType
    wagon_railway: int
    seats: List[Seat] = field(default_factory=list)

    def __str__(self) -> str:
        return f'Ticker order <From: {self.from_station_id}, To: {self.to_station_id}, ' \
               f'Wagon: {self.wagon_number}, ' \
               f'Date: {self.date}, Seats: {self.seats}>'

    @staticmethod
    def roll_random_name(name_list: List[str]) -> str:
        return random.choice(name_list)

    def add_seat(self, seat_number: str,
                 first_name: Optional[str] = None,
                 last_name: Optional[str] = None):
        if not first_name:
            first_name = self.roll_random_name(names)
        if not last_name:
            last_name = self.roll_random_name(surnames)
        seat = Seat(seat_number, first_name, last_name)
        self.seats.append(seat)
        return self

    def serialize(self) -> str:
        with open(os.path.join(os.path.abspath(os.path.dirname(__file__)),
                               'order_template.jinja2')) as fd:
            tpl = Template(fd.read())
            tpl_str = tpl.render(**vars(self))
            tpl_str = (
                quote(tpl_str.replace('\n', ''))
                    .replace('%3D', '=')
                    .replace('%26', '&')
                    .replace('%20', '')
            )
            return tpl_str


import functools
import pyopencl as cl
import numpy as np
from .sparsetensor import SparseFunction, SparseTensor
from .densetensor import GPUBuffer, DenseTensor

class GradData:
  def __init__(self, data, xidx, yidx, shape):
    self.data = data
    self.xidx = xidx
    self.yidx = yidx
    self.shape = shape

def buffer_new(ctx, shape, zero=False, dtype=np.float32):
  return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=dtype))

def buffer_np(ctx, x):
  return cl.Buffer(ctx.cl_ctx, cl.mem_flags.READ_WRITE | cl.mem_flags.COPY_HOST_PTR, hostbuf=x)

@functools.lru_cache
def clbuild(cl_ctx, name, prg):
  return cl.Program(cl_ctx, prg).build().__getattr__(name)

def uint2(x, y):
  return np.array((x,y), dtype=cl.cltypes.uint2)
i32 = np.int32

# ************* unary ops *************
def unary_op(ctx, code, x):
  ret = buffer_new(ctx, x.shape)
  unop = clbuild(ctx.cl_ctx, "unop", """
  __kernel void unop(__global const float *a_g, __global float *res_g) {
    int gid = get_global_id(0);
    float a = a_g[gid];
    res_g[gid] = """+code+""";
  }""")
  unop(ctx.cl_queue, [np.prod(ret.shape)], None, x.cl, ret.cl)
  return ret

class ReLU(SparseFunction):
  def forward(ctx, input):
    ctx.save_for_backward(input)
    return unary_op(ctx, 'max(a, (float)0.)', input)

  def backward(ctx, grad_output):
    input, = ctx.saved_tensors
    return binary_op(ctx, 'a * (b >= 0)', grad_output, input)

class Log(SparseFunction):
  def forward(ctx, input):
    ctx.save_for_backward(input)
    return unary_op(ctx, 'log(a)', input)

  def backward(ctx, grad_output):
    input, = ctx.saved_tensors
    return binary_op(ctx, 'a / b', grad_output, input)

class Exp(SparseFunction):
  def forward(ctx, input):
    ret = unary_op(ctx, 'exp(a)', input)
    ctx.save_for_backward(ret)
    return ret

  def backward(ctx, grad_output):
    ret, = ctx.saved_tensors
    return binary_op(ctx, 'a * b', grad_output, ret)

# ************* reduce ops *************

def reduce_op(ctx, code, code2, inp, axis=None, start="0.0"):
  if axis is None:
    # full reduce
    osize = [1]*len(inp.shape)
  else:
    osize = np.array(inp.shape)
    osize[list(axis)] = 1
  ret = buffer_new(ctx, osize)
  if axis is None:
    ret.shape = (1,)

  # TODO: this is insanely slow
  reduce = clbuild(ctx.cl_ctx, "reduce", """
  __kernel void reduce(__global const float *a_g, int sz, __global float *res_g, int prod, int n_dims,
                       __global const int *shape_x, __global const int *shape_ret) {
    int gid = get_global_id(0);

    float out = """+start+""";
    for (int x = 0; x < sz; x++) {
      int idx = 0;  // compute index into a_g
      int tprod = prod;
      int tsz = sz;
      for (int dim = 0; dim < n_dims; dim++) {
        idx *= shape_x[dim];
        if (shape_x[dim] == shape_ret[dim]) {   // dim from gid, don't reduce
          tprod /= shape_x[dim];
          idx += (gid / tprod) % shape_x[dim];
        } else {  // dim from x
          tsz /= shape_x[dim];
          idx += (x / tsz) % shape_x[dim];
        }
      }
      float a = a_g[idx];
      """+code+""";
    }
    res_g[gid] = """+code2+""";
  }""")
  reduce(ctx.cl_queue, [np.prod(osize)], None, inp.cl,
    i32(np.prod(inp.shape)//np.prod(osize)), ret.cl,
    i32(np.prod(osize)), i32(len(osize)),
    buffer_np(ctx, np.array(inp.shape, dtype=np.int32)),
    buffer_np(ctx, np.array(osize, dtype=np.int32)))
  return ret

class Sum(SparseFunction):
  def forward(ctx, input, axis=None):
    if isinstance(axis, int): axis = [axis]
    ctx.save_for_backward(input, axis)
    ret = reduce_op(ctx, "out += a", "out", input, axis=axis)
    if axis is not None:
      ret.shape = tuple([input.shape[i] for i in range(len(input.shape)) if i not in axis])
    return ret

  def backward(ctx, grad_output):
    input, axis = ctx.saved_tensors
    shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
    output = GPUBuffer(shape, hostbuf=grad_output)
    return binary_op(ctx, 'a+b', output, buffer_new(ctx, input.shape, zero=True))

class Max(SparseFunction):
  def forward(ctx, input, axis=None):
    if isinstance(axis, int): axis = [axis]
    ret = reduce_op(ctx, "out = max(a,out)", "out", input, axis=axis, start="-INFINITY")
    ctx.save_for_backward(input, axis, ret)
    if axis is not None:
      ret.shape = tuple([input.shape[i] for i in range(len(input.shape)) if i not in axis])
    return ret

  def backward(ctx, grad_output):
    input, axis, ret = ctx.saved_tensors
    shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
    ret2 = binary_op(ctx, "1.0*(a==b)", input, GPUBuffer(shape, ret))
    div = reduce_op(ctx, "out += a", "out+1e-10", ret2, axis=axis)
    ret3 = binary_op(ctx, "a/b", ret2, GPUBuffer(shape, div))
    return binary_op(ctx, 'a*b', ret3, GPUBuffer(shape, grad_output))

# ************* binary ops *************

@functools.lru_cache
def get_binop_prg(cl_ctx, code, complist):
  ndims = len(complist)
  args = "".join([f", int d{i}" for i in range(ndims)] + [f", int p{i}" for i in range(ndims-1)])
  compute_idx_rets = "".join([f"\n    int idx_ret{i} = (gid0 / {f'p{i}' if i < ndims-1 else '1'}) % d{i};" for i in range(ndims)])

  idx_exprs = ["0", "0"] # [idx_x, idx_y]
  for i in range(ndims):
    for j in range(2):
      if complist[i][j]:
        idx_exprs[j] = "idx_ret%d + d%d*(%s)" % (i, i, idx_exprs[j])

  return cl.Program(cl_ctx, """__kernel void binop(__global const float *x_g, __global const float *y_g, __global float *res_g"""+args+""") {
    int gid0 = get_global_id(0);"""+compute_idx_rets+"""
    float a = x_g["""+idx_exprs[0]+"""];
    float b = y_g["""+idx_exprs[1]+"""];
    res_g[gid0] = """+code+""";\n}""").build()

def binary_op(ctx, code, x, y):
  n_dims = max(len(x.shape), len(y.shape))
  shape_x, shape_y = np.ones(n_dims, dtype=np.int32), np.ones(n_dims, dtype=np.int32)
  shape_x[:len(x.shape)] = np.array(x.shape, dtype=np.int32)
  shape_y[:len(y.shape)] = np.array(y.shape, dtype=np.int32)
  if not np.all((shape_x == 1) | (shape_y == 1) | (shape_x == shape_y)):
    raise Exception(f"binary op unbroadcastable shape mismatch: {x.shape} vs {y.shape}")
  shape_ret = np.maximum(shape_x, shape_y)

  dimlist, complist = [], [] # note: len(dimlist) may be less than n_dims
  def push(dim, comp):
    if len(complist) > 0 and complist[-1] == comp:
      dimlist[-1] *= dim
    elif comp != (False, False):
      dimlist.append(dim); complist.append(comp)
  for i in range(n_dims): # group together any adjacent dimensions that we can to simplify broadcasting
    push(i32(max(shape_x[i], shape_y[i])), (shape_x[i] > 1, shape_y[i] > 1))

  prg = get_binop_prg(ctx.cl_ctx, code, tuple(complist))
  ret = buffer_new(ctx, shape_ret, zero=True)
  prod_list = np.array(dimlist, dtype=i32)[-1::-1].cumprod(dtype=i32)[-1::-1] # take cumprod from back to front
  prg.binop(ctx.cl_queue, [prod_list[0]] if len(dimlist) > 0 else [1], None, x.cl, y.cl, ret.cl, *dimlist, *(prod_list[1:]))
  return ret

def unbroadcast(ctx, out, in_sh):
  sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1] if in_sh != (1,) else None
  return reduce_op(ctx, "out += a", "out", out, sum_axis)

class Add(SparseFunction):
  def forward(ctx, x, y):
    ctx.save_for_backward(x.shape, y.shape)
    return binary_op(ctx, 'a+b', x, y)

  def backward(ctx, grad_output):
    grad_x, grad_y = grad_output, grad_output
    shape_x, shape_y = ctx.saved_tensors
    return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y),

class Sub(SparseFunction):
  def forward(ctx, x, y):
    ctx.save_for_backward(x.shape, y.shape)
    return binary_op(ctx, 'a-b', x, y)

  def backward(ctx, grad_output):
    grad_x, grad_y = grad_output, unary_op(ctx, '-a', grad_output)
    shape_x, shape_y = ctx.saved_tensors
    return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y),

class Mul(SparseFunction):
  def forward(ctx, x, y):
    ctx.save_for_backward(x, y)
    return binary_op(ctx, 'a*b', x, y)

  def backward(ctx, grad_output):
    x,y = ctx.saved_tensors
    grad_x = binary_op(ctx, 'a*b', y, grad_output)
    grad_y = binary_op(ctx, 'a*b', x, grad_output)
    return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape),

class Pow(SparseFunction):
  def forward(ctx, x, y):
    ctx.save_for_backward(x, y)
    return binary_op(ctx, 'pow(a,b)', x, y)

  def backward(ctx, grad_output):
    x,y = ctx.saved_tensors
    grad_x = binary_op(ctx, 'a*b', grad_output,
                      binary_op(ctx, 'b * (pow((float)a, (float)(b-1.0)))', x, y))
    grad_y = binary_op(ctx, 'a*b', grad_output,
                      binary_op(ctx, 'pow(a, (float)b) * log(a);', x, y))
    return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape),

# ************* movement ops *************

class Reshape(SparseFunction):
  def forward(ctx, x, shape):
    x.data.shape = tuple(shape)
    ctx.save_for_backward(x)
    return x

  def backward(ctx, grad_output):
    in_shape = ctx.saved_tensors
    return in_shape

def perm_axis(ctx, inp, order):
  # print("PERM:", inp, order)
  osize = np.array(inp.shape)[list(order)]
  ret = buffer_new(ctx, osize)
  perm = clbuild(ctx.cl_ctx, "perm", """
  __kernel void perm(__global const float *a_g, __global float *res_g, int n_axis,
                       __global const int *shape, __global const int *order) {
    int gid = get_global_id(0);
    int gi = gid;
    int idx = 0;
    for(int i = n_axis-1; i>-1; i--) {
      int stride = 1;
      for(int j=order[i]+1; j<n_axis; j++) stride *= shape[j];
      idx += (gi % shape[order[i]])*stride;
      gi /= shape[order[i]];
    }
    res_g[gid] = a_g[idx];
    }""")
  perm(ctx.cl_queue, [np.prod(osize)], None, inp.cl, ret.cl, i32(len(osize)),
    buffer_np(ctx, np.array(inp.shape, dtype=np.int32)),
    buffer_np(ctx, np.array(order, dtype=np.int32)))
  # print("RAN")
  return ret

class Transpose(SparseFunction):
  def forward(ctx, x):
    # print("T FWD:", x)
    newdata = {
     'data': x.datat,
     'idxs': x.idxst,
     'nnzs': x.nnzst,
     'ellw': x.ellwt,
     'datat': x.data,
     'idxst': x.idxs,
     'nnzst': x.nnzs,
     'ellwt': x.ellw,
    }
    newshape = tuple(np.array(x.shape).T)
    ret = SparseTensor(from_datas=newdata, shape=newshape)
    return ret

  def backward(ctx, grad_output):
    return perm_axis(ctx, grad_output, np.argsort((1,0)))

# TODO: merge this with perm axis
def inner_slice(ctx, x, arg):
  shift = [y[0] for y in arg]
  oshape = [y[1]-y[0] for y in arg]
  ret = buffer_new(ctx, oshape)
  gslice = clbuild(ctx.cl_ctx, "gslice", """
  __kernel void gslice(__global const float *input, __global float *output, int prod, int n_dims,
                       __global const int *shape_x, __global const int *shape_ret,
                       __global const int *shift) {
    int gid = get_global_id(0);
    int iptr = 0;
    int zero = 1;
    for (int dim = 0; dim < n_dims; dim++) {
      prod /= shape_ret[dim];
      int sidx = (gid / prod) % shape_ret[dim] + shift[dim];
      zero &= (sidx >= 0 && sidx < shape_x[dim]);
      iptr = (iptr * shape_x[dim]) + sidx;
    }
    output[gid] = zero ? input[iptr] : 0.0;
  }""")
  gslice(ctx.cl_queue, [np.prod(ret.shape)], None,
    x.cl, ret.cl, i32(np.prod(ret.shape)), i32(len(ret.shape)),
    buffer_np(ctx, np.array(x.shape, dtype=np.int32)),
    buffer_np(ctx, np.array(ret.shape, dtype=np.int32)),
    buffer_np(ctx, np.array(shift, dtype=np.int32)))
  return ret

class Slice(SparseFunction):
  def forward(ctx, x, arg=None):
    ctx.save_for_backward(x.shape)
    return inner_slice(ctx, x, arg)

  def backward(ctx, grad_output):
    shape, = ctx.saved_tensors
    narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(ctx.arg)]
    return inner_slice(ctx, grad_output, narg)

# ************* processing ops *************

class Matmul(SparseFunction): # input and weights are swapped, legacy..
  def forward(ctx, weight, input):
    # print("WEIGHT/input:", weight.shape, input.shape)
    # print(input.shape, weight.shape)
    # assert weight.shape[-2] == input.shape[-1]

    # if not weight.m:
    #   weight.m = DenseTensor(np.zeros((input.shape[0], weight.shape[1])))

    isize, msize, osize = i32(input.shape[-2]), i32(input.shape[-1]), i32(weight.shape[-1])
    outshape = np.array([input.shape[-2], weight.shape[-1]])
    # print("OUT:", outshape, isize, msize, osize)
    outdata = np.zeros(outshape)
    ret = DenseTensor(outdata)
    # ret = buffer_new(ctx.cl_ctx, outshape, zero=True)
    # print("RET:", ret)
    # print("RET:", input)

    matmul = clbuild(ctx.cl_ctx, "matmul", """
    // DENSE x SPARSE
    __kernel void matmul(__global  float* matData,     // INPUT MATRIX DATA
                            __global  uint*  colIdx,
                            __global  uint*  rowNnz,
                            uint   ellwidth,
                            uint   mwidth,
                            uint   ncols,
                            __global  float* vector_x,    // INPUT
                            __global  float* vector_y    // OUTPUT
                            ) { // LOCAL SHARED BUFFER
      uint gid = get_global_id(0);
      uint nrows = get_global_size(0);

      for (uint gid2 = 0; gid2 < ncols; gid2++) {
        uint nnz = rowNnz[gid2];
        float sum = 0;
        for (uint i = 0; i < nnz; i++) {
          uint index   = (gid2 * ellwidth) + i;
          uint col     = colIdx[index];
          float aval  = matData[index];
          float xval  = vector_x[gid*mwidth+col];
          sum  += aval * xval;
          //if (gid==0 && gid2==0)
          //  printf("aval, xval: %.2f,%.2f - %.2f: (%i,%i) \\n", aval, xval, sum, col, index);
        }
        //printf("SUM/NNZ: %.2f %i \\n", sum, nnz);
        vector_y[gid*ncols+gid2] = sum;
      }
    }""")
    ctx.save_for_backward(input, weight)

    # (isize,msize) x (msize,osize) = (isize,osize)
    matmul(ctx.cl_queue, [outshape.T[0]], None,
      weight.datat.cl, weight.idxst.cl, weight.nnzst.cl, np.uint32(weight.ellwt), np.uint32(msize), np.uint32(outshape.T[1]), input.cl, ret.data.cl)

    # resa = np.zeros(isize,osize).astype(np.float32)
    # cl.enqueue_copy(ctx.cl_queue, resa, ret.cl)
    # return ret.data
    # return trans_axis(ctx, ret.data, (1,0))   # print("RES:", resa)
    return ret.data

  def backward(ctx, grad_output):
    input, weight = ctx.saved_tensors
    topkx, topky = weight.topkx, weight.topky
    # print('BACK:', weight.shape, topkx, topky)
    isize, msize, osize = i32(input.shape[-2]), i32(input.shape[-1]), i32(weight.shape[-1])

    grad_input = DenseTensor(np.zeros(input.shape), dtype=np.float32)
    grad_weight = DenseTensor(np.zeros(weight.shape), dtype=np.float32)

    # print('GO:', input.shape, grad_output.shape)
    # print("OUTSHAPE:", weight.shape, input.shape[0], isize, msize, weight.ellwt)

    # grad_output = grad_output + weight.m

    matmul2 = clbuild(ctx.cl_ctx, "matmul2", """
    // DENSE x SPARSE-T
    __kernel void matmul2(__global  float* matData,     // INPUT MATRIX DATA
                            __global  uint*  colIdx,
                            __global  uint*  rowNnz,
                            uint   ellwidth,
                            uint   mwidth,
                            uint   ncols0,
                            __global  float* vector_x,    // INPUT
                            __global  float* vector_y    // OUTPUT
                            ) { // LOCAL SHARED BUFFER
      uint gid = get_global_id(0);
      uint nrows = get_global_size(0);
      uint nnz = rowNnz[gid];
      uint gid2 = get_global_id(1);
      uint ncols = get_global_size(1);

      float sum = 0;
      for (uint i = 0; i < nnz; i++) {
        uint index   = (gid2 * ellwidth) + i;
        uint col     = colIdx[index];
        float aval  = matData[index];
        float xval  = vector_x[gid*mwidth+col];
        sum  += aval * xval;
        //if (gid==1 && gid2==0) {
        //  printf("aval, xval: %.2f,%.2f - %.2f: (%i,%i) \\n", aval, xval, sum, col, index);
        //}
      }
      //printf("SUM/NNZ: %.2f %i \\n", sum, nnz);
      vector_y[gid*ncols+gid2] = sum;
    }""")
    # (isize,osize) x (msize,osize) = (isize,msize)
    # print('msize:', grad_output.shape, input.shape)
    matmul2(ctx.cl_queue, input.shape, None,
      weight.data.cl, weight.idxs.cl, weight.nnzs.cl, np.uint32(weight.ellw), np.uint32(grad_output.shape[1]), np.uint32(input.shape[0]), grad_output.cl, grad_input.data.cl)

    # resa = np.zeros((input.shape[1], input.shape[0])).astype(np.float32)
    # cl.enqueue_copy(ctx.cl_queue, resa, grad_input.data.cl)

    # print('INPUT', DenseTensor(input).cpu().data, weight.shape[0], weight.shape[1])
    # print('OUT:', grad_input.cpu().data)

    gettopkx = clbuild(ctx.cl_ctx, "gettopkx", """
    // multilplies x TRANSPOSED by y (dense-dense)
    __kernel void gettopkx(__global  float* x,      // INPUT MATRIX DATA
                          __global  float* xsum,    // INPUT
                          __global  uint*  youtidx, // OUT
                          uint topky,
                          uint msize
                          ) { // LOCAL SHARED BUFFER
      uint isize = get_global_size(0);
      int gidx = get_global_id(0); // row

      // get topk
      xsum[gidx] = 0;
      for (uint i=0; i<msize; i++) {
        float val = x[i*isize+gidx];
        //if (gid == 0) {
        //  printf("\\nADD VALx: %.2f - %i", val, i*msize+gid);
        //}
        xsum[gidx] += val;
      }

      float valx = xsum[gidx];
      uint posx = 0;
      for (uint i = 0; i < isize; i++) {
        float tempval = fabs(xsum[i]);
        bool larger = (tempval > fabs(valx)) || (fabs(tempval) == fabs(valx) && i < gidx);
        posx += (larger)?1:0;
      }
      if (posx < topky) {
        youtidx[posx] = gidx;
      }
    }""")

    gettopky = clbuild(ctx.cl_ctx, "gettopky", """
    // multilplies x TRANSPOSED by y (dense-dense)
    __kernel void gettopky(__global  float* y,      // INPUT
                          __global  float* ysum,    // INPUT
                          __global  uint*  xoutidx, // OUT
                          uint topkx,
                          uint msize
                          ) { // LOCAL SHARED BUFFER
      uint osize = get_global_size(0);
      int gidy = get_global_id(0); // row

      ysum[gidy] = 0;
      for (uint i=0; i<msize; i++) {
        float val = y[i*osize+gidy];
        ysum[gidy] += val;
      }
      //barrier(CLK_GLOBAL_MEM_FENCE);
      float valy = ysum[gidy];
      uint posy = 0;
      for (uint i = 0; i < osize; i++) {
        float tempval = fabs(ysum[i]);
        bool larger = (tempval > fabs(valy)) || (fabs(tempval) == fabs(valy) && i < gidy);
        posy += (larger)?1:0;
      }
      if (posy < topkx) {
        xoutidx[posy] = gidy;
      }
    }""")

    sortuints = clbuild(ctx.cl_ctx, "sortuints", """
    // multilplies x TRANSPOSED by y (dense-dense)
    __kernel void sortuints(__global  uint* x,      // INPUT MATRIX DATA
                            __global  uint* xs      // INPUT
                            ) { // LOCAL SHARED BUFFER
      uint isize = get_global_size(0);
      int gidx = get_global_id(0); // row

      uint val = x[gidx];
      uint posx = 0;
      for (uint i = 0; i < isize; i++) {
        uint tempval = x[i];
        bool smaller = tempval < val;
        posx += (smaller)?1:0;
      }
      xs[posx] = x[gidx];
    }""")

    matmul0 = clbuild(ctx.cl_ctx, "matmul0", """
    // multilplies x TRANSPOSED by y (dense-dense)
    __kernel void matmul0(__global  float* x,      // INPUT MATRIX DATA
                          __global  float* y,      // INPUT
                          __global  uint* xidx,   // INPUT YIDX
                          __global  uint* yidx,   // INPUT YIDX
                          __global  float* resdata,// OUT
                          __global  uint*  rescols,
                          __global  uint*  resnnzs,
                          uint topkx,
                          uint ellw,
                          uint isize,
                          uint msize,
                          uint osize
                          ) { // LOCAL SHARED BUFFER

      uint topky = get_global_size(0);
      uint gidx = yidx[get_global_id(0)]; // row
      for (uint gidy0 = 0; gidy0 < topkx; gidy0++) {
        uint gidy = xidx[gidy0];
        float ret = 0.0;
        uint i;
        for (i = 0; i < msize; i++) {
          uint xidx = i*isize+gidx;
          float xval = x[xidx];
          uint yidx = osize*i+gidy;
          float yval = y[yidx];
          ret += xval*yval;
          //if (gidx==0 && gidy==0)
          //  printf("\\nmult: %.2f x %.2f - %.2f  -- %i/%i", xval, yval, ret, xidx, yidx);
        }
        //if (gidx==0&&gidy==0)
        //  printf("\\nsum:%.2f", ret);

        // add for
        uint nnz = resnnzs[gidx];
        for (i = 0; i < nnz; i++) {
          if (rescols[i] >= gidy) {
            break;
          }
          for (uint j = nnz; j >= i; j--) {
            //resdata[j+1] = resdata[j];
          }
        }
        resdata[gidx * ellw + gidy0] = ret;
        rescols[gidx * ellw + gidy0] = gidy;
        resnnzs[gidx] += 1;
      }
    }""")

    matmul0t = clbuild(ctx.cl_ctx, "matmul0t", """
    // multilplies x TRANSPOSED by y (dense-dense)
    __kernel void matmul0t(__global  float* x,      // INPUT MATRIX DATA
                          __global  float* y,      // INPUT
                          __global  uint* xidx,   // INPUT YIDX
                          __global  uint* yidx,   // INPUT YIDX
                          __global  float* resdata,// OUT
                          __global  uint*  rescols,
                          __global  uint*  resnnzs,
                          uint topky,
                          uint ellw,
                          uint isize,
                          uint msize,
                          uint osize
                          ) { // LOCAL SHARED BUFFER
      uint topkx = get_global_size(0);
      uint gidy = xidx[get_global_id(0)]; // row
      for (uint gidx0 = 0; gidx0 < topky; gidx0++) {
        uint gidx = yidx[gidx0];
        float ret = 0.0;
        uint i;
        for (i = 0; i < msize; i++) {
          uint xidx = i*isize+gidx;
          float xval = x[xidx];
          uint yidx = osize*i+gidy;
          float yval = y[yidx];
          ret += xval*yval;
          //if (gidx==0 && gidy==0)
          //  printf("\\nmult: %.2f x %.2f - %.2f  -- %i/%i", xval, yval, ret, gidx, gidy,i);
        }
        //if (gidx==0&&gidy==0)
        //  printf("\\nsum:%.2f", ret);

        // add for
        uint nnz = resnnzs[gidx];
        for (i = 0; i < nnz; i++) {
          if (rescols[i] >= gidy) {
            break;
          }
          for (uint j = nnz; j >= i; j--) {
            //resdata[j+1] = resdata[j];
          }
        }
        resdata[gidy * ellw + gidx0] = ret;
        rescols[gidy * ellw + gidx0] = gidx;
        resnnzs[gidy] += 1;
      }
    }""")

    # Weight update
    isize = weight.shape[0]
    msize = grad_output.shape[0]
    osize = weight.shape[1]

    dim1 = weight.shape[1]#min(weight.shape[1], topkx)
    dim2 = weight.shape[0]#min(weight.shape[0], topky)

    x_sum_buf   = DenseTensor(np.zeros(weight.shape[0]))
    y_sum_buf   = DenseTensor(np.zeros(weight.shape[1]))
    x_idx_buf   = DenseTensor(np.zeros(topkx), dtype=np.uint32)
    y_idx_buf   = DenseTensor(np.zeros(topky), dtype=np.uint32)
    xs_idx_buf  = DenseTensor(np.zeros(topkx), dtype=np.uint32)
    ys_idx_buf  = DenseTensor(np.zeros(topky), dtype=np.uint32)
    sdata_buf   = DenseTensor(np.zeros(weight.shape[0]*topkx))
    sidxs_buf   = DenseTensor(np.zeros(weight.shape[0]*topkx), dtype=np.uint32)
    snnzs_buf   = DenseTensor(np.zeros(weight.shape[0]), dtype=np.uint32)
    sdatat_buf  = DenseTensor(np.zeros(weight.shape[1]*topky))
    sidxst_buf  = DenseTensor(np.zeros(weight.shape[1]*topky), dtype=np.uint32)
    snnzst_buf  = DenseTensor(np.zeros(weight.shape[1]), dtype=np.uint32)

    # print('IN', DenseTensor(input).cpu().data, weight.shape, input.shape[0])
    # print('INPUT', grad_output.cpu().data)
    # print('OUT', grad_input.cpu().data.sum())

    # print('asdf:', isize, msize, osize)
    gettopkx(ctx.cl_queue,  [isize], None, input.cl, x_sum_buf.data.cl,
      y_idx_buf.data.cl, np.uint32(topky), np.uint32(msize))
    gettopky(ctx.cl_queue,  [osize], None, grad_output.cl,
      y_sum_buf.data.cl, x_idx_buf.data.cl, np.uint32(topkx), np.uint32(msize))
    sortuints(ctx.cl_queue,  [topkx], None, x_idx_buf.data.cl, xs_idx_buf.data.cl)
    sortuints(ctx.cl_queue,  [topky], None, y_idx_buf.data.cl, ys_idx_buf.data.cl)

    matmul0(ctx.cl_queue,  [topky], None, input.cl, grad_output.cl, xs_idx_buf.data.cl,
      ys_idx_buf.data.cl, sdata_buf.data.cl, sidxs_buf.data.cl, snnzs_buf.data.cl,
      np.uint32(topkx), np.uint32(topkx), np.uint32(isize), np.uint32(msize), np.uint32(osize))
    matmul0t(ctx.cl_queue, [topkx], None, input.cl, grad_output.cl, xs_idx_buf.data.cl,
      ys_idx_buf.data.cl, sdatat_buf.data.cl, sidxst_buf.data.cl, snnzst_buf.data.cl,
      np.uint32(topky), np.uint32(topky), np.uint32(isize), np.uint32(msize), np.uint32(osize))

    # x_sum_buf.data.cl.release()
    # y_sum_buf.data.cl.release()
    # sdata_buf.data.cl.release()
    # sidxs_buf.data.cl.release()
    # snnzs_buf.data.cl.release()
    # sdatat_buf.data.cl.release()
    # sidxst_buf.data.cl.release()
    # snnzst_buf.data.cl.release()
    # x_idx_buf.data.cl.release()
    # y_idx_buf.data.cl.release()

    newdata = {
     'data': sdata_buf.data,
     'idxs': sidxs_buf.data,
     'nnzs': snnzs_buf.data,
     'ellw': topkx,
     'datat': sdatat_buf.data,
     'idxst': sidxst_buf.data,
     'nnzst': snnzst_buf.data,
     'ellwt': topky,
    }

    w_grad = SparseTensor(from_datas=newdata, shape=weight.shape)
    # gradpy = w_grad.to_numpy()
    # print('grad_max:', w_grad.shape, gradpy.sum())
    # gradpy = w_grad.to_numpy(dual=True)
    # print('grad_max:', w_grad.shape, gradpy.sum())
    # asdf

    # updatem = clbuild(ctx.cl_ctx, "updatem", """
    # // sorts x and y in ascending order and returns sorted indices
    # __kernel void updatem(__global  float* m,     // INPUT MATRIX DATA
    #                       __global  float* grad,     // INPUT MATRIX DATA
    #                       uint msize,
    #                       uint osize,
    #                       uint topk,
    #                       float scale,
    #                       __global  uint*  xoutidx,
    #                       __global  uint*  youtidx,
    #                       __global  float* matData,     // OUTPUT MATRIX DATA
    #                       __global  uint*  colIdx,
    #                       __global  uint*  rowNnz
    #                       ) {
    #   uint gid = get_global_id(0);
    #   uint nnz = rowNnz[gid];
    #   for (uint i=0; i<nnz; i++) {
    #     uint col = colIdx[gid*topk+i];
    #     float val = matData[gid*topk+i];
    #     m[osize*gid+col] = 0;
    #   }
    #   for (uint i=0; i<osize; i++) {
    #     m[osize*gid+i] = scale * grad[osize*gid+i];
    #   }
    # }""")
    # scale = 0.9
    # updatem(ctx.cl_queue, [grad_output.shape[0],], None,
    #   weight.m.data.cl, grad_output.data.cl, np.uint32(grad_input.shape[-1]), np.uint32(grad_output.shape[1]), np.uint32(topky), np.float32(scale), xs_idx_buf.data.cl, ys_idx_buf.data.cl,
    #   sdata_buf.data.cl, sidxs_buf.data.cl, snnzs_buf.data.cl)

    return w_grad, grad_input

def trans_axis(ctx, inp, order=(1,0)):
  osize = np.array(inp.shape)[list(order)]
  ret = buffer_new(ctx, osize)
  trans = clbuild(ctx.cl_ctx, "trans", """
    __kernel void trans(__global float *a_g,
                        __global float *res_g,
                        uint width) {
      int row = get_global_id(0);
      for(uint i=0; i<width; i++) {
        //printf("\\nSET:%i-%i", row, i);
        res_g[row*width+i] = 0;
      }
    }""")
  trans(ctx.cl_queue, [osize[1]], None, inp.cl, ret.cl, np.uint32(osize[0]))

  print("PERM RET:", ret)
  return ret

class Conv2D(SparseFunction):
  def forward(ctx, x, w, stride=1, groups=1):
    if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
    cout,cin,H,W = w.shape
    ys,xs = ctx.stride
    bs,cin_,iy,ix = x.shape
    oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
    if cin*ctx.groups != cin_: raise Exception(f"Input Tensor shape {x.shape} does not match the shape of the weights {w.shape}. ({cin*ctx.groups} vs. {cin_})")
    assert cout % ctx.groups == 0
    rcout = cout//ctx.groups

    ctx.save_for_backward(x,w)

    # output buffer
    ret = buffer_new(ctx, (bs, cout, oy, ox))

    # input  = (bs, groups, cin, iy, ix)
    # weight = (groups, rcout, cin, H, W)
    # output = (bs, groups, rcout, oy, ox)

    conv = clbuild(ctx.cl_ctx, "conv", """
    __kernel void conv(__global const float *input, __global const float *weight, __global float *output,
      int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs) {

      int B = get_global_id(0)/(groups*rcout);  // range 0-bs
      int g = (get_global_id(0)/rcout)%groups;
      int c = get_global_id(0) % rcout;

      int Y = get_global_id(1);  // range 0-oy
      int X = get_global_id(2);  // range 0-ox
      int IY = Y*ys;
      int IX = X*xs;

      float acc = 0.0;
      for (int ci = 0; ci < cin; ci++) {
        for (int y = IY; y < IY+H; y++) {
          for (int x = IX; x < IX+W; x++) {
            acc += input[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + y*ix + x] * \
              weight[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + (y-IY)*W + (x-IX)];
          }
        }
      }
      output[B*groups*rcout*oy*ox + g*rcout*oy*ox + c*oy*ox + Y*ox + X] = acc;
    }""")

    conv(ctx.cl_queue, [bs*groups*rcout, oy, ox], None,
      x.cl, w.cl, ret.cl,
      i32(H), i32(W), i32(groups), i32(rcout), i32(cin),
      i32(oy), i32(ox), i32(iy), i32(ix), i32(ys), i32(xs)
    )
    return ret

  def backward(ctx, grad_output):
    bs,_,oy,ox = grad_output.shape
    x, w = ctx.saved_tensors
    cout,cin,H,W = w.shape
    ys,xs = ctx.stride
    bs,cin_,iy,ix = x.shape
    oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
    assert cin*ctx.groups == cin_
    assert cout % ctx.groups == 0
    rcout = cout//ctx.groups

    dx = buffer_new(ctx, (bs, cin_, iy, ix), zero=True)
    dw = buffer_new(ctx, (cout, cin, H, W))

    # tensx = (bs, groups*cin, iy, ix)
    # tensw = (groups*rcout, cin, H, W)
    # ggg = (bs, groups*rout, oy, ox)

    convw = clbuild(ctx.cl_ctx, "convw", """
    __kernel void convw(__global const float *tensx, __global const float *ggg, __global float *dw,
      int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {

      int g = get_global_id(0)/(rcout*cin) ; // range 0-groups
      int c = (get_global_id(0)/(cin)) %rcout; // range 0-rcout
      int ci = get_global_id(0) % cin;        // range 0-cin
      int y = get_global_id(1);  // range 0-H
      int x = get_global_id(2);  // range 0-W

      float acc = 0.0;
      for (int Y = 0; Y < oy; Y++) {
        for (int X = 0; X < ox; X++) {
          for (int B = 0; B < bs; B++) {
            acc += ggg[B*groups*rcout*oy*ox + +g*rcout*oy*ox + c*oy*ox + Y*ox + X] * \
              tensx[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + (Y*ys+y)*ix + X*xs+x];
          }
        }
      }
      dw[get_global_id(0)*H*W + y*W + x] = acc;
    }""")
    convx = clbuild(ctx.cl_ctx, "convx", """
    __kernel void convx(__global const float *tensw, __global const float *ggg, __global float *dx,
      int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {

      int B = get_global_id(0);
      int g = get_global_id(1);
      int ci = get_global_id(2);

      for (int Y = 0; Y < oy; Y++) {
        for (int X = 0; X < ox; X++) {
          for (int y = 0; y < H; y++) {
            for (int x = 0; x < W; x++) {
              float acc = 0.0;
              for (int c = 0; c < rcout; c++) {
                acc += ggg[B*groups*rcout*oy*ox + g*rcout*oy*ox + c*oy*ox + Y*ox + X] * \
                  tensw[g*rcout*cin*H*W + c*cin*H*W + ci*H*W + y*W + x];
              }
              dx[B*groups*cin*iy*ix + g*cin*iy*ix + ci*iy*ix + (Y*ys+y)*ix + X*xs+x] += acc;
            }
          }
        }
      }
    }
    """)

    conv_args = i32(H), i32(W), i32(ctx.groups), i32(rcout), i32(cin), i32(oy), i32(ox), i32(iy), i32(ix), i32(ys), i32(xs), i32(bs)
    convw(ctx.cl_queue, [ctx.groups*rcout*cin, H, W], None, x.cl, grad_output.cl, dw.cl, *conv_args)
    convx(ctx.cl_queue, [bs, ctx.groups, cin], None, w.cl, grad_output.cl, dx.cl, *conv_args)
    return dx, dw


def bfs(graph, vis, node, q):
    while q:
        sz = len(q)
        for i in range(sz):
            node = q.pop(0)
            if vis[node] == True:
                continue
            print(node, end = " ")
            for adj in graph[node]:
                q.append(adj)
        print()

n = int(input("Enter number of nodes: "))
m = int(input("Enter number of edges: "))

graph = [[] for i in range(n)]

print("Enter nodes in form of u, v")

for count in range(m):
    u, v = map(int, input().strip().split())
    graph[u].append(v)

q = []
q.append(0)
vis = [False for i in range(n)]
node = 0

print()
bfs(graph, vis, node, q)
print()


r'''
    Python library for standard functions required for numerical methods.
    - Method f() - returns value of polynomial at given value.
    - Method dof() - returns value of derivative of polynomial at given value.
    - Method value() - returns value of expression at given value of x.
'''
import lambdaFunction
def f(x: int | float, equ_list: list, deg: int) -> int | float :
    '''
    Function for obtaining value of any polynomial at given value of x.

    Args :
        x : value of x
        equ_list = list of coeffs of x in decreasing order.
        deg : degree of polynomial.

    Returns:
        returns float or int.
    '''
    power = deg
    sum = 0
    for coeff in equ_list:
        sum = sum + (coeff * (x**power))
        power = power - 1
    return sum
def dof(x: int | float, equ_list: list, deg: int) -> int | float:
    '''
    Function for obtaining value of derivative of any polynomial at given value of x.

    Args :
        x : value of x
        equ_list = list of coeffs of x in decreasing order.
        deg : degree of polynomial.

    Returns:
        returns float or int.
    '''
    power = deg
    diff = [0]*power
    i = 0
    j = power
    while(i<power):
        new_coeff = equ_list[i] * j
        diff[i] = new_coeff
        j = j - 1
        i = i + 1
    return f(x, diff, (len(diff)-1))

def value(exp : str, x : int | float) -> int | float:
    '''
    Function for obtaining value of given expression at given value of x.

    Args :
        exp : Expression
        x : value of x

    Returns :
        returns float or int.
    '''
    return lambdaFunction.lambdafunction("def f(x):", f"\n\treturn {exp}", caller=f"f({x})")


from ete3 import Tree , TreeStyle , NodeStyle , faces , AttrFace
import csv

#t = Tree("(A:1,(B:1,(E:1,D:1):0.5):0.5);" )
#t.render("mytree.png", w=183, units="mm")

t = Tree("(0);" )

all_orgs = []
with open('testlineage.csv') as csvfile:
    reader = csv.DictReader(csvfile)
    for org in reader:
        all_orgs.append((org['id'],org['ancestor_id'],org['recorded_at'],org['on_lod']))
with open('testunpruned.csv') as csvfile:
    reader = csv.DictReader(csvfile)
    for org in reader:
        all_orgs.append((org['id'],org['ancestor_id'],'101','0'))

all_orgs.sort(key=lambda t: int(t[0]))
#print (all_orgs)

style = NodeStyle()
style["fgcolor"] = "#00f000"
#style["bgcolor"] = "lightred"
#style["size"] = 10
style["shape"] = "circle"
style["vt_line_color"] = "#0000aa"
style["hz_line_color"] = "#0000aa"
style["vt_line_width"] = 50
style["hz_line_width"] = 50
style["vt_line_type"] = 2 # 0 solid, 1 dashed, 2 dotted
style["hz_line_type"] = 2

for org in all_orgs:
    #print (org)
    new_node = Tree('(' + org[0]  + ');') #+ ':' + str(int(org[2])^1)
    if (org[3] == '1'):
        #print (org)
        new_node.set_style(style)
    anc = t.search_nodes(name=org[1])[0]
    anc.add_child(new_node)

#for n in t.traverse():
    #print (n.name)
    #if all_orgs[int(n.name)][3] == '1':
        #n.img_style = style2

style2 = NodeStyle()
style2["fgcolor"] = "#000000"
style2["shape"] = "circle"
style2["vt_line_color"] = "#0000aa"
style2["hz_line_color"] = "#0000aa"
style2["vt_line_width"] = 5
style2["hz_line_width"] = 5
style2["vt_line_type"] = 1 # 0 solid, 1 dashed, 2 dotted
style2["hz_line_type"] = 1
for l in t.iter_leaves():
    l.img_style = style2

ts = TreeStyle()
ts.show_leaf_name = True
#ts.rotation = 90
ts.mode = 'c'
ts.arc_start = 180
ts.arc_span = 350
#t.show(tree_style=ts)
#t.show()

t = Tree()
t.populate(8)
style2 = NodeStyle()
style2["fgcolor"] = "darkred"
style2["shape"] = "circle"
style2["vt_line_color"] = "green"
style2["hz_line_color"] = "red"
style2["vt_line_width"] = 5
style2["hz_line_width"] = 5
style2["vt_line_type"] = 1 # 0 solid, 1 dashed, 2 dotted
style2["hz_line_type"] = 1
for l in t.iter_leaves():
    l.img_style = style2
ts = TreeStyle()
ts.show_leaf_name = False
#
t.show(tree_style=ts)


"""
@author: nabin

This script runs pdb2seq.pl to get sequence from pdb file

"""
import os
import sys

def run_pdb2sec(pdb, density_name, input_path, perl_script_dir):
    try:
        density_map_dir = os.path.join(input_path, density_name)
        pdb_fi = os.path.join(density_map_dir, pdb)
        os.system("perl " + perl_script_dir + " " + pdb_fi + ">>" + density_map_dir + "/" + "atomic.fasta")
        print(density_name, "Done")
    except FileNotFoundError:
        print("Error for file:", density_name)

if __name__ == "__main__":
    input_path = sys.argv[1]
    perl_script_dir = "pdb2seq.pl"
    density_maps = [den for den in os.listdir(input_path) if os.path.isdir(os.path.join(input_path, den))]
    for den in density_maps:
        rm1 = f'{input_path}/{den}/atomic.fasta'
        if os.path.exists(rm1):
            os.remove(rm1)
        pdb_file = [p for p in os.listdir(os.path.join(input_path, den)) if
                    p.endswith(".pdb") or p.endswith(".ent")]
        pdb_file.sort()
        pdb_file_n = pdb_file[0].split(".")[0]
        pdb_file_na = pdb_file_n.split("_")[0]
        pdb_file_name = pdb_file_na + ".pdb"
        # print(den, "->", pdb_file_name)

        run_pdb2sec(pdb_file_name, den, input_path, perl_script_dir)


'''class is user defined class and used to access more informations'''
class A:
    param1= "ras"
    param2="rashu"
    def fun(self):
        print(f'i am {self.param1}')
        print(f'i am {self.param2}')

obj= A()
print(obj.param1)
obj.fun()


import tflearn as tl
import numpy as np
import os,glob,cv2
import sys,argparse
from read_image import read_valid_image
import tensorflow as tf
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.core import input_data , dropout, fully_connected
# from tflearn.layers.estimator import  regression

image_data,image_label,num_class,label_name = read_valid_image()
# First, pass the path of the image
dir_path = os.path.dirname(os.path.realpath(__file__))
image_path=sys.argv[1]
filename = dir_path +'/' +image_path
image_size=32
num_channels=1
images = []
# Reading the image using OpenCV
image = cv2.imread(filename,0)

# Resizing the image to our desired size and preprocessing will be done exactly as done during training
image = cv2.resize(image, (image_size, image_size), cv2.INTER_LINEAR)
image = image[:, :, np.newaxis]
images.append(image)
images = np.array(images, dtype=np.uint8)
images = images.astype('float32')
x = np.multiply(images, 1.0/255.0)
# y = np.zeros([1,2])
#The input to the network is of shape [None image_size image_size num_channels]. Hence we reshape.
# x_batch = images.reshape(1, image_size,image_size,num_channels)
convnet = input_data(shape=[None,32,32,1],name='input')

convnet = conv_2d(convnet,32,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = conv_2d(convnet,64,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = conv_2d(convnet,32,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = conv_2d(convnet,64,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = conv_2d(convnet,32,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = conv_2d(convnet,64,2,activation='relu')
convnet = max_pool_2d(convnet,2)

convnet = fully_connected(convnet, 1024,activation='relu')

convnet = fully_connected(convnet,2,activation='softmax',name="output")
model = tl.DNN(convnet)
model.load('tflearn')
final = model.predict(x)
for i in range(0,len(label_name)):
    print(label_name[i],final[0][i])

# print(final)


from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import absolute_import
from builtins import *  # NOQA
from future import standard_library
standard_library.install_aliases()

import chainer
from chainer import functions as F
import numpy as np

import chainerrl

def _wrap_by_variable(x):
    if isinstance(x, chainer.Variable):
        return x
    else:
        return chainer.Variable(x)

def _unwrap_variable(x):
    if isinstance(x, chainer.Variable):
        return x.data
    else:
        return x

def elementwise_gaussian_log_pdf(x, mean, var, ln_var):
    # log N(x|mean,var)
    #   = -0.5log(2pi) - 0.5log(var) - (x - mean)**2 / (2*var)
    return -0.5 * np.log(2 * np.pi) - \
        0.5 * ln_var - \
        ((x - mean) ** 2) / (2 * var)

NPY_SQRT1_2 = 1 / (2 ** 0.5)

def _ndtr(a):
    """CDF of the standard normal distribution.

    See https://github.com/scipy/scipy/blob/master/scipy/special/cephes/ndtr.c
    """
    if not isinstance(a, chainer.Variable):
        a = chainer.Variable(a)
    x = a * NPY_SQRT1_2
    z = abs(x)
    half_erfc_z = 0.5 * F.erfc(z)
    return F.where(
        z.data < NPY_SQRT1_2,
        0.5 + 0.5 * F.erf(x),
        F.where(
            x.data > 0,
            1.0 - half_erfc_z,
            half_erfc_z))

def _safe_log(x):
    """Logarithm function that won't backprop inf to input."""
    return F.log(F.where(x.data > 0, x, x.data))

def _log_ndtr(x):
    """Log CDF of the standard normal distribution.

    See https://github.com/scipy/scipy/blob/master/scipy/special/cephes/ndtr.c
    """
    if not isinstance(x, chainer.Variable):
        x = chainer.Variable(x)
    return F.where(
        x.data > 6,
        -_ndtr(-x),
        F.where(
            x.data > -14,
            _safe_log(_ndtr(x)),
            -0.5 * x * x - _safe_log(-x) - 0.5 * np.log(2 * np.pi)))

def _gaussian_log_cdf(x, mu, sigma):
    """Log CDF of a normal distribution."""
    if not isinstance(x, chainer.Variable):
        x = chainer.Variable(x)
    return _log_ndtr((x - mu) / sigma)

def _gaussian_log_sf(x, mu, sigma):
    """Log SF of a normal distribution."""
    if not isinstance(x, chainer.Variable):
        x = chainer.Variable(x)
    return _log_ndtr(-(x - mu) / sigma)

class ClippedGaussian(chainerrl.distribution.GaussianDistribution):
    """Clipped Gaussian distribution."""

    def __init__(self, mean, var, low, high):
        super().__init__(mean, var)
        self.low = F.broadcast_to(low, mean.shape)
        self.high = F.broadcast_to(high, mean.shape)
        assert isinstance(self.low, chainer.Variable)
        assert isinstance(self.high, chainer.Variable)

    def sample(self):
        unclipped = F.gaussian(self.mean, self.ln_var)
        return F.minimum(F.maximum(unclipped, self.low), self.high)

    def log_prob(self, x):
        unclipped_elementwise_log_prob = elementwise_gaussian_log_pdf(
            x, self.mean, self.var, self.ln_var)
        std = self.var ** 0.5
        low_log_prob = _gaussian_log_cdf(self.low, self.mean, std)
        high_log_prob = _gaussian_log_sf(self.high, self.mean, std)
        x_data = _unwrap_variable(x)
        elementwise_log_prob = F.where(
            (x_data <= self.low.data),
            low_log_prob,
            F.where(
                x_data >= self.high.data,
                high_log_prob,
                unclipped_elementwise_log_prob))
        return F.sum(elementwise_log_prob, axis=1)

    def prob(self, x):
        return F.exp(self.log_prob(x))

    def copy(self):
        return ClippedGaussian(_wrap_by_variable(self.mean.data),
                               _wrap_by_variable(self.var.data),
                               self.low,
                               self.high)