import numpy
import scipy.interpolate
from scipy.ndimage import filters
import scipy.stats.mstats
import logging

logger = logging.getLogger(__name__)
formatter = logging.Formatter(
    "%(asctime)s - %(name)s - %(levelname)s - %(message)s")

def lat_weights(lats):
    """ Returns cos(lats)

    Required as a weighting for calculating areas on
    a regular cartesian lat-lon grid in which meridians
    converge.

    Arguments:
       lats (numpy ndarray): in degree

    Returns:
       numpy ndarray
    """
    wgts = numpy.cos(numpy.radians(lats))
    return wgts

def runave(data, N, axis=0, step=None):
    """ Compute a smooth box-car running average and masked
    edge values using scipy.ndimage.filters.convolve1d

    Args:
       data (numpy array)
       N (int) : width of the box car
       axis (int): axis along which the average is computed
                (See scipy.ndimage.filters.convolve1d)
       step (optional, int): skips between the box car samples

    Returns:
       numpy masked array (same shape as data)
    """
    if isinstance(data, numpy.ma.core.MaskedArray):
        data = data.filled(numpy.nan)
    if step is None:
        weights = numpy.ones((N,))
    else:
        if type(step) is not int:
            raise Exception("step should be an integer")
        weights = numpy.array(([1.]+[0.]*(step-1))*(N-1)+[1.])
    weights /= weights.sum()
    return numpy.ma.masked_invalid(filters.convolve1d(
        data, weights, mode='constant', cval=numpy.nan, axis=axis))

def skewness(data):
    """ Essentially the same as scipy.stats.skew or
        scipy.stats.mstats.skew depending on the input

    Input:
        data (numpy array or masked array)
    Returns:
        numpy array or masked array
    """
    if isinstance(data,numpy.ma.core.MaskedArray):
        npmod = numpy.ma
    else:
        npmod = numpy
    return npmod.power(data-npmod.mean(data), 3).mean()/\
        npmod.power(npmod.var(data), 3./2.)

def resample_xy(x, y, xnew, nx, ny, y_always_valid=True):
    """ Given paired values of (x,y), randomly sample a set of
    values for ynew given an array xnew such that the joint distribution
    of (xnew, ynew) resembles that of (x,y)

    Arguments:
        x (numpy 1d array)
        y (numpy 1d array): shape should match x
        xnew (numpy 1d array)
        nx (int) : the number of bins applied to x
        ny (int) : the number of bins applied to y
        y_always_valid (bool): keep drawingn from the cdf until y is valid

    Returns:
        ynew (numpy 1d array): length = len(xnew)
    """
    if y.shape != x.shape:
        raise ValueError("shape of y should match the shape of x")

    xedges = numpy.linspace(x.min(), x.max(), nx+1)
    xnew_ibin = numpy.digitize(xnew, xedges)
    y_cdf = {}
    y_mids = {}
    ynew = numpy.empty_like(xnew, dtype=y.dtype)
    for ix, ibin in enumerate(xnew_ibin):
        # If the xnew is out of bound, return numpy.nan for that entry
        if xnew[ix] >= xedges[-1] or xnew[ix] < xedges[0]:
            logger.warn("Out of bound for x={:.2e}. "+\
                        "Return numpy.nan.".format(float(xnew[ix])))
            ynew[ix] = numpy.nan
            continue

        # Don't compute the cdf twice for the same bin
        # Compute it once and save it
        if ibin not in y_cdf:
            ys = y[(x > xedges[ibin-1]) & (x <= xedges[ibin])]
            y_cdf[ibin], y_mids[ibin] = cdf(ys,bins=ny)

        # Keep drawing a new y until the drawn value is valid
        while True:
            rand_num = numpy.random.uniform()
            try:
                ynew[ix] = scipy.interpolate.interp1d(
                    y_cdf[ibin], y_mids[ibin])(rand_num)
                break
            except ValueError:
                if y_always_valid:
                    continue
                else:
                    ynew[ix] = numpy.nan
                    break
    return ynew

def cdf(data, **kwargs):
    """ Compute a histogram and the corresponding cumulative
    frequency polygon.

    Args:
        data (numpy 1d array)
        option (str): currently accepts "hist" only

    Keyword arguments currently are passed to numpy.histogram

    Returns:
        bins, cdf : mid values of bins and the cumulative frequencies
    """
    # Use histogram for cdf
    h, x = numpy.histogram(data,**kwargs)
    x_mid = 0.5*(x[1:]+x[:-1])
    h_cum = numpy.cumsum(h)
    return h_cum/float(h_cum[-1]), x_mid


import argparse

from vinfraclient.cmd.base import ShowOne

class ShowSettings(ShowOne):
    _description = 'Show NVMe performance settings.'

    def do_action(self, parsed_args):
        return self.app.vinfra.multiple_cses.settings.show_params()

class ChangeSettings(ShowOne):
    _description = 'Enabling/Disabling NVMe performance.'

    __known_tiers = 0, 1, 2, 3

    @classmethod
    def number_of_cses_per_disk(cls, number):
        try:
            number = int(number)
        except ValueError:
            raise argparse.ArgumentTypeError('1 or more CSes per disk must be specified')

        if number < 1:
            raise argparse.ArgumentTypeError("Minimum number of CSes per disk is 1")

        return number

    def configure_parser(self, parser):
        for tier in self.__known_tiers:
            parser.add_argument(
                '--tier%d' % tier,
                type=self.number_of_cses_per_disk,
                metavar='number',
                default=None,
                help='Set number of CSes per disk for tier %d' % tier,
            )

    def do_action(self, parsed_args):
        kwargs = {}
        for tier in self.__known_tiers:
            number = getattr(parsed_args, 'tier%d' % tier)
            if number is None:
                continue
            kwargs['tier%d' % tier] = number

        return self.app.vinfra.multiple_cses.settings.change_params(**kwargs)


from typing import Any

class QueueEmpty(IndexError):
    pass

class InvalidFixedLength(ValueError):
    pass

class Queue(object):

    """A queue (FIFO) using a set of fixed-length `arrays`.

    Supports `enqueue`, `dequeue`, and `get_size` operations.
    """

    def __init__(self, fixed_length: int):
        """
        :param fixed_length: Length of internal fixed-length arrays
        """

        try:
            fixed_length = int(fixed_length)
        except ValueError:
            raise InvalidFixedLength(f"Queue could not covert {fixed_length} to valid fixed-length")
        except TypeError:
            raise InvalidFixedLength(f"Queue does not accept {type(fixed_length)}")

        if fixed_length <= 0:
            raise InvalidFixedLength(f"Fixed-length must be greater than 0, got: {fixed_length}")

        self._arrays = []
        self._fixed_length = fixed_length

        # Index to enqueue on
        self._enqueue_idx = fixed_length

        # Index to dequeue from
        self._dequeue_idx = 0

        # Number of elements
        self._count = 0

    @property
    def fixed_length(self):
        """Fixed length used by arrays"""
        return self._fixed_length

    @property
    def n_arrays(self):
        """Number of fixed length arrays in-use"""
        return len(self._arrays)

    def enqueue(self, item: Any):
        """Puts an object onto end of queue"""

        # Allocate new array if needed
        if self._enqueue_idx >= self._fixed_length:
            self._enqueue_idx = 0
            self._arrays.append([None] * self._fixed_length)

        # Enqueue data
        self._arrays[-1][self._enqueue_idx] = item
        self._enqueue_idx += 1
        self._count += 1

    def dequeue(self) -> Any:
        """Pops the object from start of queue"""

        # Size guard
        if self._count <= 0:
            raise QueueEmpty("Cannot dequeue while empty")

        # Remove dequeue array if exhausted
        if self._dequeue_idx >= self._fixed_length:
            self._dequeue_idx = 0
            self._arrays = self._arrays[1:]

        # Get first item
        item = self._arrays[0][self._dequeue_idx]
        self._dequeue_idx += 1
        self._count -= 1

        # Reset original array if empty
        if self._count == 0:
            self._enqueue_idx = 0
            self._dequeue_idx = 0
            self._arrays = self._arrays[:1]

        return item

    def get_size(self) -> int:
        """Returns current size of queue"""
        return self._count


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

from odoo import models, fields, api
from odoo.exceptions import ValidationError, UserError
from datetime import datetime, timedelta
from odoo.tools import float_is_zero, float_compare, DEFAULT_SERVER_DATETIME_FORMAT

class B2bSaleResendWizard(models.TransientModel):
    _name = 'b2b.sale.resend.wizard'
    _description = 're-send products claimed by end customer'

    origin = fields.Many2one('sale.order', u'销售单号')
    order_line = fields.One2many('b2b.sale.resend.wizard.line', 'order_id', u'订单明细')
    amount_total = fields.Float(u'合计金额', compute='_get_total_amount')
    # procurement_ids = fields.One2many('procurement.order', 'sale_line_id', string='Procurements')
    procurement_group_id = fields.Many2one('procurement.group', 'Procurement Group', copy=False)

    @api.onchange('origin')
    def onchange_origin(self):
        freight_product = self.env.ref('b2b_platform.default_b2b_delivery_carrier')
        list = []
        if self.origin:
            for line in self.origin.order_line:
                if line.product_id != freight_product:
                    list.append([0,0,{'order_id': self.id,
                                      'product_id': line.product_id,
                                      'qty': line.product_uom_qty,
                                      'price_unit': line.price_unit,
                                      'sale_line': line.id,
                                    }])
            self.order_line = list

    @api.one
    def _get_total_amount(self):
        amount = 0
        for line in self.order_line:
            amount += line.qty * line.price_unit
        self.amount_total = amount

    def _prepare_procurement_group(self):
        return {'name': self.origin.name}

    @api.multi
    def btn_confirm(self):
        if not self.order_line:
            raise UserError(u'没有需要采购的产品')

        p_amount = 0
        freight_product = self.env.ref('b2b_platform.default_b2b_delivery_carrier')
        freight_obj = self.env['b2b.product.freight.line'].sudo()

        for p_line in self.order_line:
            p_amount += p_line.price_unit * p_line.qty

        if p_amount > self.origin.partner_id.deposit_avail_amt:
            raise UserError(u'在平台预存的余额不足，请先充值')

        # 计算运费
        if not self.origin.partner_shipping_id.country_id:
            raise UserError(u'发货地址未明确所在国家，无法核算运费')
        dest_country = self.origin.partner_shipping_id.country_id
        amount = 0
        freight_line_id = False
        for line in self.order_line:
            if line.qty <= 0:
                line.unlink()
            else:
                freight = freight_obj.search([('product_id', '=', line.product_id.product_tmpl_id.id),
                                              ('country_id', '=', dest_country.id)], limit=1)
                if not freight:
                    raise UserError(u'未找到%s产品的运费' % line.product_id.name)
                amount += freight.cny_amount * line.qty
                if line.product_id == freight_product:
                    freight_line_id = line
        if freight_line_id:
            freight_line_id.price_unit = amount
        else:
            self.order_line.create({
                'order_id': self.id,
                'price_unit': amount,
                'qty': 1,
                'product_id': freight_product.id,
                'is_delivery': True,
            })

        proc_order_line = self.order_line.search([('order_id', '=', self.id), ('is_delivery', '!=', True)])

        # 创建发货单
        purchase_obj = self.env['purchase.order'].sudo()
        picking_obj = self.env['stock.picking']
        warehouse_obj = self.env['stock.warehouse'].sudo()
        warehouse = warehouse_obj.search([('partner_id', '=', 1)], limit=1)
        if not warehouse:
            raise UserError(u'未找到平台仓库编码，请咨询平台管理人员')

        stock_loc_obj = self.env['stock.location'].sudo()
        cust_loc_id = stock_loc_obj.search([('usage', '=', 'customer')], limit=1)
        if not cust_loc_id:
            raise UserError(u'未找到客户所属库位，请咨询平台管理人员')

        picking_type_obj = self.env['stock.picking.type'].sudo()
        picking_type = picking_type_obj.search([('id', '=', self.env.ref('stock.picking_type_out').id), ('warehouse_id', '=', warehouse.id)], limit=1)
        # picking_type = picking_type_obj.search([('name', '=', u'发货'), ('warehouse_id', '=', warehouse.id)], limit=1)
        if not picking_type:
            raise UserError(u'未找到捡货类型，请咨询平台管理人员')

        # 生成procurement order
        new_procs = self.env['procurement.order']  # Empty recordset
        for line in proc_order_line:
            if line.product_id != freight_product:
                new_proc = new_procs.search([('origin','=',self.origin.name),
                                             ('sale_line_id','=',line.sale_line),
                                             ('state','=','done')],limit=1).copy()
                new_proc.write({'state':'confirmed', 'product_qty':line.qty, 'product_uom_qty':line.qty, 'ordered_qty':line.qty})
                new_procs += new_proc
        # 如果平台有库存，则从平台自有仓库发货：判断是否需从平台发货
        purchase = False
        platform = False
        for line in proc_order_line:
            if line.qty > line.product_id.own_stock:
                purchase = True
                line.is_procure = True
            else:
                platform = True
                line.is_procure = False
        # 如果需从平台库存发货
        if platform:
            # 生成拣货单草稿
            picking = picking_obj.sudo().create({'location_id': self.env.ref('b2b_platform.stock_location_wh_own').id,
                                          'picking_type_id': picking_type.id,
                                          'location_dest_id': cust_loc_id.id,
                                          'ship_address': self.origin.partner_shipping_id.id,
                                          'state': 'draft',
                                          'move_type': 'one',
                                          # 'carrier_id': self.carrier_id.id,
                                          'origin': self.origin.name,
                                          'partner_id': self.origin.partner_id.id,
                                          'owner_id': 1,
                                          'company_id_id': 1,
                                          'min_date': datetime.today(),
                                          })
            for pl_line in proc_order_line:
                if not pl_line.is_procure:
                    move_line = self.env['stock.move'].search([('origin', '=', self.origin.name),
                                                               ('product_id', '=', pl_line.product_id.id),
                                                               ('product_qty', '=', pl_line.qty),
                                                               ('state', '!=', 'done'),
                                                               ], limit=1)
                    if move_line:
                        move_line.sudo().write({'picking_id': picking.id})
        # 如果需从经销商采购
        if purchase:
            # 拆分采购单
            supp_list = proc_order_line.mapped('prod_supplier')
            for supplier in supp_list:
                supp_loc_id = stock_loc_obj.search([('location_id.name', '=', u'供应商库存'),
                                                    ('partner_id', '=', supplier.id)], limit=1)
                if not supp_loc_id:
                    raise UserError(u'未找到供应商所属库位，请咨询平台管理人员')

                # 生成拣货单草稿
                picking = picking_obj.sudo().create({'location_id': supp_loc_id.id,
                                              'picking_type_id': picking_type.id,
                                              'location_dest_id': cust_loc_id.id,
                                              'ship_address': self.origin.partner_shipping_id.id,
                                              'state': 'draft',
                                              'move_type': 'one',
                                              # 'carrier_id': self.carrier_id.id,
                                              'origin': self.origin.name,
                                              'partner_id': self.origin.partner_id.id,
                                              'owner_id': supplier.id,
                                              'company_id_id': 1,
                                              'min_date': datetime.today(),
                                              })
                move_lines = self.env['stock.move'].sudo().search([('origin', '=', self.origin.name),
                                                                    ('product_id.product_owner', '=', supplier.id),
                                                                    ('state', '!=', 'done')])
                if move_lines:
                    move_lines.write({'state': 'draft'})
                    move_lines.write({'picking_id': picking.id})
                    for m_line in move_lines:
                        ord_qty = m_line.procurement_id.product_qty
                        m_line.write({'product_uom_qty': ord_qty,'ordered_qty': ord_qty})

                # 根据库存情况决定供货路线
                # 如果供应商是经销商自己，则不需要生成采购单
                if supplier != self.origin.shop_id.parent_id:
                    supp_order_line = self.order_line.search([('order_id', '=', self.id),
                                                              ('need_procure', '=', True),
                                                              ('is_procure', '=', True),
                                                              ('is_delivery', '!=', True),
                                                              ('prod_supplier', '=', supplier.id)])
                    if supp_order_line:
                        # 创建并确认采购订单
                        amount = 0
                        po_list = []
                        for line in supp_order_line:
                            if line.product_id != freight_product:
                                freight = freight_obj.search([('product_id', '=', line.product_id.product_tmpl_id.id),
                                                              ('country_id', '=', dest_country.id)], limit=1)
                                amount += freight.cny_amount * line.qty
                                po_list.append([0, 0, {'product_id': line.product_id.id,
                                                       'name': line.product_id.name,
                                                       'date_planned': datetime.today(),
                                                       'product_qty': line.qty,
                                                       'product_uom': line.product_id.product_tmpl_id.uom_id.id,
                                                       'price_unit': line.price_unit,
                                                       'taxs_id': False,
                                                       }])
                        if amount:
                            po_list.append([0, 0, {'product_id': freight_product.id,
                                                   'name': freight_product.name,
                                                   'date_planned': datetime.today(),
                                                   'product_qty': 1,
                                                   'product_uom': freight_product.id,
                                                   'price_unit': amount,
                                                   'taxs_id': False,
                                                   }])
                        purchase_id = purchase_obj.create({
                            'partner_id': supplier.id,
                            'currency_id': self.env.user.company_id.currency_id.id,
                            'origin': self.origin.name,
                            'origin_so': self.origin.id,
                            'distributor': self.origin.partner_id.id,
                            'date_order': datetime.today(),
                            'date_planned': datetime.today(),
                            'picking_type_id': self.env['stock.picking.type'].sudo().search([('name', '=', u'直运')],
                                                                                limit=1).id,
                            'dest_address_id': self.origin.partner_shipping_id.id,
                            'order_line': po_list
                        })
                        purchase_id.state = 'purchase'

        # 创建并确认客户发票
        cust_inv_obj = self.env['account.invoice'].sudo()
        cust_inv_line_obj = self.env['account.invoice.line'].sudo()
        cust_inv_id = cust_inv_obj.create({
            'journal_id': self.env['account.journal'].sudo().search([('type', '=', 'sale')],
                                                                    limit=1).id or 1,
            'type': 'out_invoice',
            'partner_id': self.origin.partner_id.id,
            'reference': self.origin.name,
            'origin': self.origin.name,
            'amount_total_company_signed': self.amount_total,
            'residual': self.amount_total,
            'amount_untaxed': self.amount_total,
            'residual_company_signed': self.amount_total,
            'amount_total_signed': self.amount_total,
            'residual_signed': self.amount_total,
            'amount_total': self.amount_total,
            'amount_untaxed_signed': self.amount_total,
            'state': 'draft',
            'account_id': 4,
            'currency_id': 8,
            'name': self.origin.name,
            'commercial_partner_id': self.origin.partner_id.id,
            # 'partner_invoice_id': self.partner_invoice_id.id,
            'partner_shipping_id': self.origin.partner_id.id,
        })
        for line in self.order_line:
            cust_inv_line_obj.create({
                'origin': self.origin.name,
                'price_unit': line.price_unit,
                'price_subtotal': line.amount_subtotal,
                'currency_id': self.env.user.company_id.currency_id.id,
                'uom_id': line.product_id.product_tmpl_id.uom_id.id,
                'partner_id': self.origin.partner_id.id,
                'company_id': 1,
                'account_id': 52,
                'price_subtotal_signed': line.amount_subtotal,
                'name': u'100% 的采购货款',
                'product_id': line.product_id.id,
                'invoice_id': cust_inv_id.id,
                'quantity': line.qty,
                })
            line.write({'invoice_status': 'invoiced'})
        cust_inv_id.action_invoice_open()

        self.origin.write({'qdoo_state': 'part-deliver', 'invoice_status': 'invoiced'})

        # return {'name': u'产品清单',
        #         'type': 'ir.actions.act_window',
        #         'res_model': 'product.product',
        #         'view_type': 'form',
        #         'view_mode': 'tree',
        #         # 'view_id': 'view_tree_b2b_distributor_product_list',
        #         # 'search_view_id': 'search_b2b_distributor_product_list',
        #         'domain': [('my_products', '=', True)],
        #         'context': {'create': False},
        #         }

class B2bSaleResendWizardLine(models.TransientModel):
    _name = 'b2b.sale.resend.wizard.line'

    order_id = fields.Many2one('b2b.sale.resend.wizard', u'订单号')
    product_id = fields.Many2one('product.product', u'产品')
    qty = fields.Float(u'数量')
    price_unit = fields.Float(u'单价')
    amount_subtotal = fields.Float(u'金额', compute='_get_subtotal')
    need_procure = fields.Boolean(u'走平台', default=True)
    is_procure = fields.Boolean(u'走供应商')
    own_product = fields.Boolean(u'自有', compute="_if_own_product")
    prod_supplier = fields.Many2one('res.partner', u'供应商', related='product_id.product_owner', readonly=True)
    sale_line = fields.Integer(u'订单行ID')
    is_delivery = fields.Boolean(u'运费')

    @api.one
    def _get_subtotal(self):
        self.amount_subtotal = self.qty * self.price_unit

    @api.multi
    def _prepare_order_line_procurement(self, group_id=False):
        self.ensure_one()
        return {
            'name': self.product_id.name,
            'origin': self.order_id.origin.name,
            'date_planned': datetime.strptime(self.order_id.origin.date_order, DEFAULT_SERVER_DATETIME_FORMAT),
            'product_id': self.product_id.id,
            'product_qty': self.qty,
            'product_uom': self.product_id.product_tmpl_id.uom_id.id,
            'company_id': self.order_id.origin.company_id.id,
            'group_id': group_id,
            'sale_line_id': self.sale_line,
        }


from unified.core.auth import Auth
from datetime import datetime, timezone
import requests
import json

TYPEFORM_BASE_URL = 'https://api.typeform.com'

def rest(method, url, body=None, token=None):

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }

    return requests.request(method, url, headers=headers, json=body)

def epoch_to_format(format, epoch):
    """Convert  epoch to given format"""

    return datetime.fromtimestamp(int(epoch[:10]), tz=timezone.utc).strftime(format)

def get_typeform_answer(typeform_answers, field_id):
    """Get the answer for the field with given"""

    for answer in typeform_answers:
        if field_id == answer['field']['id']:
            if answer['type'] == 'choice':
                return answer[answer['type']]['label']
            else:
                return answer[answer['type']]

    return None

def remove_field(pl, key):
    """Remove items with given key in any dict at any level in the nested dict/list/tuple"""

    if type(pl) is tuple:
        r = (remove_field(v, key) for v in pl)

    elif type(pl) is list:
        r = [remove_field(v, key) for v in pl]

    elif type(pl) is dict:
        r = {k: remove_field(v, key) for (k, v) in pl.items() if k != key}
    else:
        r = pl

    return r


from selenium import webdriver
import time
from selenium.webdriver.common.keys import Keys
from bs4 import BeautifulSoup
import csv
import os
import openpyxl
from tkinter import Tk
import argparse
from selenium.webdriver.common.keys import Keys
from webdriver_manager.chrome import ChromeDriverManager
import re

def scrap_from_url(args, driver,url=None, load_page=True):
    page = 0
    max_page = 55
    extracted_infos = []
    count = 0
    item_num = 1
    user_ids = set()
    while page < max_page:
        links = []
        page += 1
        print('page: ', page)
        names = driver.find_elements_by_tag_name("a")
        for name in names:
            link = name.get_attribute('href')
            if 'https://shopping.yahoo.co.jp/review/item/list?store_id=' in link:
                links.append(link)
        # print('num of link: ', len(links))
        for link in links:
            count += 1
            print(count)
            url=link
            url_splt = re.split('[?.,=_/&]',url)
            user_id = url_splt[url_splt.index('id') + 1]

            if user_id not in user_ids:
                 user_ids.add(user_id)
                 print('Totally', len(user_ids), 'stores have been scanned')
            else:
                continue
            template1 = 'https://store.shopping.yahoo.co.jp/'
            template2 = '/info.html'
            info_page = template1 + user_id + template2

            driver.get(info_page)
            if ('error' in driver.current_url) or ('notfound' in driver.current_url):
                continue
            info = driver.find_element_by_tag_name('body').text.split()
            person_profile = {}
            shop_name = info[info.index('会社名（商号）') + 1]
            address = info[info.index('住所') + 1]
            manager_name = info[info.index('代表者') + 1]
            phone=''
            fax=''
            email=''
            for pos,element in enumerate(info):
                if 'お問い合わせ' == element:
                    phone = info[pos + 1]
                    email = info[pos + 2]
                    break
                if 'お問い合わせ電話番号' == element:
                    phone = info[pos + 1]
                if 'お問い合わせファックス番号' == element:
                    fax = info[pos + 1]
                if 'お問い合わせメールアドレス' == element:
                    email = info[pos + 1]
                    break
            person_profile['SHOP NAME'] = shop_name
            person_profile['MANAGER'] = manager_name
            person_profile['ADDRESS'] = address
            person_profile['PHONE'] = phone
            person_profile['FAX'] = fax
            person_profile['EMAIL'] = email
            extracted_infos.append(person_profile)
        if page % 2 == 0:
            with open('results_yahoo14.csv', 'w', encoding='utf8', newline='') as output_file:
                fc = csv.DictWriter(output_file, fieldnames=extracted_infos[0].keys())
                fc.writeheader()
                fc.writerows(extracted_infos)
        item_num += 30
        next_page = 'https://shopping.yahoo.co.jp/category/13457/2494/37524/list?X=8&view=list&sc_i=shp_pc_catelist_nrwcgt&b='
        next_page = next_page + str(item_num)
        driver.get(next_page)
    with open('results_yahoo14.csv', 'w', encoding='utf8', newline='') as output_file:
        fc = csv.DictWriter(output_file, fieldnames=extracted_infos[0].keys())
        fc.writeheader()
        fc.writerows(extracted_infos)
    print('Done')
####################################################################################
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--search_words', default=None)
    parser.add_argument('--url', default=None)
    parser.add_argument('--mode', default='url')
    args = parser.parse_args()

    # driver = webdriver.Chrome(ChromeDriverManager(version="91.0.4472.101").install())
    driver = webdriver.Chrome()
    # driver.set_page_load_timeout(1)
    driver.get('https://shopping.yahoo.co.jp/category/13457/2494/37524/list?X=8&view=list&sc_i=shp_pc_catelist_nrwcgt&b=1')
    # driver.execute_script("driver.stop();")
    # driver.maximize_window()
    scrap_from_url(args, driver, args.url, load_page=False)


from abc import ABC, abstractmethod
from sqlalchemy.ext.asyncio import AsyncSession
from typing import List, Optional
from sqlalchemy.orm import selectinload
from sqlalchemy import Boolean, select, and_, func

from src.database.models import *
from src.database.repositories import GenericSqlRepository, GenericRepository
from src.database.exceptions import *
from src.user.exceptions import *

class IStoreRepository(GenericRepository[Store], ABC):
    @abstractmethod
    async def get_all(
        self, limit: int | None, offset: int | None, substr: str | None
    ) -> List[Store]:
        raise NotImplementedError()

    @abstractmethod
    async def get_by_id(self, store_id: int) -> Store:
        raise NotImplementedError()

class StoreRepository(GenericSqlRepository[Store], IStoreRepository):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Store)

    async def get_all(
        self,
        limit: int | None,
        offset: int | None,
        substr: str | None,
        tag_ids: list[int] | None,
    ) -> List[Store]:
        stmt = self._construct_statement_get_all().order_by(Store.name)
        stmt = stmt.options(
            selectinload(Store.products).selectinload(Product.product_feedbacks),
            selectinload(Store.store_tag_links).selectinload(StoreTagLink.tag),
        )
        stmt = self._add_limit_offset_to_stmt(stmt, limit, offset)
        stmt = self._add_substr_to_stmt(stmt, Store.name, substr)
        if tag_ids:
            for tag in tag_ids:
                sub = self._constructor_subquery(StoreTagLink, [], tag_id=tag)
                stmt = stmt.join(sub)
        return await self._execute_statement_get_all(stmt)

    async def get_by_id(self, store_id: int) -> Store:
        stmt = self._construct_statement_get_by_id(store_id)
        stmt = stmt.options(
            selectinload(Store.products).selectinload(Product.product_feedbacks)
        )
        return await self._execute_statement_get_by_id(stmt, store_id)

class IProductRepository(GenericRepository[Product], ABC):
    @abstractmethod
    async def get_all(
        self,
        store_id: int,
        limit: int | None,
        offset: int | None,
        substr: str | None,
        tag_ids: list[int] | None,
        is_service: bool,
    ) -> List[Product]:
        raise NotImplementedError()

    @abstractmethod
    async def get_with_org(self, id: int) -> Optional[Product]:
        raise NotImplementedError()

    @abstractmethod
    async def get_by_id(self, product_id: int) -> Optional[Product]:
        raise NotImplementedError()

class ProductRepository(GenericSqlRepository[Product], IProductRepository):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Product)

    async def get_all(
        self,
        store_id: int,
        limit: int | None,
        offset: int | None,
        substr: str | None,
        tag_ids: list[int] | None,
        is_service: bool,
    ) -> List[Product]:
        stmt = self._construct_statement_get_all(store_id=store_id).order_by(
            Product.name
        )
        stmt = stmt.options(
            selectinload(Product.product_items),
            selectinload(Product.product_feedbacks),
            selectinload(Product.product_tag_links).selectinload(ProductTagLink.tag),
        )
        stmt = self._add_limit_offset_to_stmt(stmt, limit, offset)
        stmt = self._add_substr_to_stmt(stmt, Product.name, substr)
        if tag_ids:
            for tag in tag_ids:
                sub = self._constructor_subquery(ProductTagLink, [], tag_id=tag)
                stmt = stmt.join(sub)
        if is_service is not None:
            stmt = stmt.filter(Product.is_service == is_service)
        return await self._execute_statement_get_all(stmt)

    async def get_with_org(self, id: int) -> Optional[Product]:
        stmt = self._construct_statement_get_by_id(id)
        stmt = stmt.options(selectinload(Product.store))
        return await self._execute_statement_get_by_id(stmt, id)

    async def get_by_id(self, product_id: int) -> Product | None:
        stmt = self._construct_statement_get_by_id(product_id)
        stmt = stmt.options(
            selectinload(Product.product_items), selectinload(Product.product_feedbacks)
        )
        return await self._execute_statement_get_by_id(stmt, product_id)

class IProductItemRepository(GenericRepository[ProductItem], ABC):
    @abstractmethod
    async def get_all(
        self, product_id: int, limit: int | None, offset: int | None, substr: str | None
    ):
        raise NotImplementedError()

    @abstractmethod
    async def get_owner_by_id(self, product_item_id: int):
        raise NotImplementedError()

class ProductItemRepository(GenericSqlRepository[ProductItem], IProductItemRepository):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, ProductItem)

    async def get_all(
        self, product_id: int, limit: int | None, offset: int | None, substr: str | None
    ):
        stmt = self._construct_statement_get_all(offset, limit, product_id=product_id)
        stmt = self._add_substr_to_stmt(stmt, ProductItem.name, substr)
        res = await self._execute_statement_get_all(stmt)
        return res

    async def get_owner_by_id(self, product_item_id: int):
        stmt = self._construct_statement_get_by_id(product_item_id)
        stmt = stmt.options(
            selectinload(ProductItem.product)
            .selectinload(Product.store)
            .selectinload(Store.owner)
        )
        return await self._execute_statement_get_by_id(stmt, product_item_id)

class ICharacteristicRepository(GenericRepository[Characteristic], ABC):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Characteristic)

    @abstractmethod
    async def get_all(
        self, limit: int | None, offset: int | None, substr: str | None
    ) -> List[Characteristic]:
        raise NotImplementedError()

    @abstractmethod
    async def _check_exist_char(self, char: Characteristic) -> bool:
        raise NotImplementedError()

class CharacteristicRepository(
    GenericSqlRepository[Characteristic], ICharacteristicRepository
):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Characteristic)

    async def get_all(
        self, limit: int | None, offset: int | None, substr: str | None
    ) -> List[Characteristic]:
        stmt = self._construct_statement_get_all()
        stmt = self._add_limit_offset_to_stmt(stmt, limit, offset).order_by(
            Characteristic.name
        )
        stmt = self._add_substr_to_stmt(stmt, Characteristic.name, substr)
        return await self._execute_statement_get_all(stmt)

    async def _check_exist_char(self, char: Characteristic) -> bool:
        stmt = await self._session.execute(
            select(Characteristic).where(Characteristic.name == char.name)
        )
        result = stmt.scalar_one_or_none()
        return True if result else None

class IProductFeedbackRepository(GenericRepository[ProductFeedback], ABC):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, ProductFeedback)

    @abstractmethod
    async def get_by_user_product(self, user_id: int, product_id: int):
        raise NotImplementedError()

    @abstractmethod
    async def get_all(
        self, product_id: int, limit: int | None, offset: int | None, count: int | None
    ):
        pass

class ProductFeedbackRepository(
    GenericSqlRepository[ProductFeedback], IProductFeedbackRepository
):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, ProductFeedback)

    async def get_by_user_product(self, user_id: int, product_id: int):
        stmt = await self._session.execute(
            select(ProductFeedback).filter(
                ProductFeedback.user_id == user_id,
                ProductFeedback.product_id == product_id,
            )
        )
        result = stmt.scalar_one_or_none()
        return result

    async def get_all(
        self, product_id: int, limit: int | None, offset: int | None, count: int | None
    ):
        stmt = self._construct_statement_get_all(offset, limit, product_id=product_id)
        stmt = stmt.options(selectinload(ProductFeedback.likes))
        stmt = stmt.options(
            selectinload(ProductFeedback.likes),
            selectinload(ProductFeedback.user).selectinload(User.role),
        ).order_by(ProductFeedback.mark.desc())
        if count:
            sub = (
                select(ProductFeedback)
                .join(Like)
                .group_by(ProductFeedback.id)
                .having(func.count(Like.product_feedback_id) > count)
                .subquery()
            )
            stmt = stmt.select_from(sub)
        res = await self._execute_statement_get_all(stmt)
        return res

class ILikeRepository(GenericRepository[Like], ABC):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Like)

    @abstractmethod
    async def get_by_id(self, user_id: int, product_feedback_id: int):
        raise NotImplementedError()

    @abstractmethod
    async def delete(self, user_id: int, product_feedback_id: int):
        raise NotImplementedError()

class LikeRepository(GenericSqlRepository[Like], ILikeRepository):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session, Like)

    async def get_by_id(self, user_id: int, product_feedback_id: int):
        stmt = await self._session.execute(
            select(Like).filter(
                Like.user_id == user_id,
                Like.product_feedback_id == product_feedback_id,
            )
        )
        result = stmt.scalar_one_or_none()
        return result

    async def delete(self, user_id: int, product_feedback_id: int):
        stmt = select(Like).filter(
            Like.user_id == user_id,
            Like.product_feedback_id == product_feedback_id,
        )
        res = await self._execute_statement_get_one(stmt)
        await self._session.delete(res)
        await self._session.flush()


# ----------------------------------------------------------------------------
# Created By  : Tobias Mink, Marvin Kemper
# ---------------------------------------------------------------------------
"""
Creates a legend which holds metadata corresponding to the colors of the mesh.
"""
# ---------------------------------------------------------------------------
from data.dictionarys import original_layers
from data.lists import colors

def do(self, task: str):
    if task == 'remove':
        self.plotter.remove_legend()
    if task == 'add':
        new_list = []
        for mesh_name, color in zip(original_layers.keys(), colors):
            new_list.append([mesh_name, color])
        self.plotter.add_legend(labels=new_list, bcolor='#0c1726', face="r", loc="upper left", size=(0.1, 0.1))


#!/usr/bin/python3
""" Base class module """

import json

class Base:
    """ My Base class """

    __nb_objects = 0

    def __init__(self, id=None):
        ''' Instantiates the new Base object '''
        if id is not None:
            self.id = id
        else:
            Base.__nb_objects += 1
            self.id = Base.__nb_objects

    @staticmethod
    def to_json_string(list_dictionaries):
        ''' that returns the JSON string representation
        of list_dictionaries

        Args:
            list_dictionaries: is a list of dictionaries '''
        if list_dictionaries is None or list_dictionaries == '[]':
            return "[]"
        else:
            return json.dumps(list_dictionaries)

    @classmethod
    def save_to_file(cls, list_objs):
        ''' writes the JSON string representation
        of list_objs to a file '''
        filename = cls.__name__ + '.json'
        dics = []
        #lists = []
        #with open(filename, mode='w') as file:
        #    if list_objs is None:
        #        file.write('[]')
        #    else:
        #        for i in range(len(list_objs)):
        #            list_dic.append(list_objs[i].to_dictionary())

        if list_objs is None:
            return []
        for lists in list_objs:
            dics.append(lists.to_dictionary())

        with open(filename, mode='w') as file:
            file.write(Base.to_json_string(dics))

    def from_json_string(json_string):
        ''' from json to dic '''
        if json_string is None or json_string == '[]':
            return '[]'
        return json.loads(json_string)


# %%

from functools import partial
import os
import itertools
from copy import copy, deepcopy
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
from jax.interpreters.ad import JVPTracer

from tabulate import tabulate
import msgpack
import numpy as np
import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt
import plotly.express as px
from attrs import evolve, define, Factory, frozen
import einops

from interp.model.gpt_model import Gpt, partial_gpt_call
from interp.model.grad_modify_fake import ablation_softmax_probs, integrated_gradients_softmax_probs
from interp.model.grad_modify_mul import attn_score_mul_builder, probs_v_mul_builder
from interp.tools.attribution_backend_utils import stop_qkv_except
from interp.tools.grad_modify_query_items import ItemConf, ItemIdx, MulConf, NoneConf, StopGradConf, ModifyInBasisConf
from interp.tools.grad_modify_query_utils import MulBuilder, compose_trees as ct
from interp.tools.interpretability_tools import (
    batch_tokenize,
    losses_runner,
    losses_runner_log,
    print_max_min_by_tok_k,
    single_tokenize,
    LossesRunnerTreeConfig,
    begin_token,
    losses_runner_tree,
    toks_to_string_list,
    run_on_tokens,
    cross_entropy_ablation_log_probs,
    get_cross_entropy_integrated_gradients_log_probs,
)
from interp.tools.jax_util import stack_tree
from interp.model.model_loading import load_model
from interp.tools.grad_modify_query import (
    GradModifierConf,
    ModifierCollectionTreeNodeStack,
    ModifierCollectionTreeNode,
    Query,
    TargetConf,
    run_queries,
    run_query,
)
from interp.tools.indexer import I
from interp.model.grad_modify_output import (
    Attn,
    AttnLayer,
    Embeds,
    FinalOutput,
    OutputConf,
    losses_runner_outputs,
    output_tree,
)
from interp.tools.log import Idxs, KeyIdxs
from interp.ui.attribution_backend import AttributionBackend, AttributionRoot
from interp.tools.data_loading import DATA_DIR
import interp
import interp.tools.optional as op

# os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"

# %%

"""
This notebook does some quick ablations on the induction circuit to see what's
actually important.
"""

# %%

# interp.init()  # if you want interp to work, run in jupyter notebook classic

# %%

models_dir_local = os.path.expanduser("~/interp_models_jax")
model, params, tok = load_model("jan5_attn_only_two_layers/", models_dir=models_dir_local)
model: Gpt = model
params

# %%

text = " \"I don't sleep right,\" Harry said. He waved his hands helplessly. \"My sleep cycle is twenty-six hours long, I always go to sleep two hours later, every day. I can't fall asleep any earlier than that, and then the next day I go to sleep two hours later than that. 10PM, 12AM, 2AM, 4AM, until it goes around the clock. Even if I try to wake up early, it makes no difference and I'm a wreck that whole day. That's why I haven't been going to a normal school up until now.\""
text = " \"I love Mrs. Dursley. I don't sleep right,\" Harry said. He waved his hands helplessly. \"Mrs. Dursley's sleep cycle is twenty-six hours long, I always go to sleep two hours later, every day. I can't fall asleep any earlier than that, and then the next day I go to sleep two hours later than that. 10PM, 12AM, 2AM, 4AM, until it goes around the clock. Mr. Dursley stops me from sleeping two hours everytime. Mr. Dursley dislikes his wife Mrs. Dursley."
# text = " My name is Buck. My name is Buck."
example_seq_no_begin = batch_tokenize([text])
example_seq = jnp.concatenate([jnp.expand_dims(begin_token(), (0, 1)), example_seq_no_begin], axis=1)
example_seq_string_list = toks_to_string_list(example_seq[0])
example_seq_names = [f"{n} ({i})" for i, n in enumerate(example_seq_string_list)]
example_seq.shape

# %%

import torch

fnames = os.listdir(f"{DATA_DIR}/owt_tokens_int16_val")[1:2]
all_tokens = [torch.load(f"{DATA_DIR}/owt_tokens_int16/{f}") for f in fnames]
data_pt = list(itertools.chain(*[torch.split(x["tokens"], x["lens"].tolist()) for x in all_tokens]))

# %%

max_size = 511

# %%

data = torch.stack(
    [data_pt_val[:max_size].to(torch.int64) + 32768 for data_pt_val in data_pt if data_pt_val.size(0) >= max_size],
    dim=0,
).numpy()
data.shape

# %%

data_subset = data[:300]
batch_size = 4

# %%

"""
We'll want to ablate 'through' the log softmax instead of differentiating.

We'll also see what the derivative and integrated derivative look like.
"""

# %%

out_tree = output_tree(model, OutputConf(AttnLayer(1, by_head=True), endpoint=FinalOutput(), item_idx=ItemIdx(I[:, 4])))

extra_tree: List[List[Optional[ModifierCollectionTreeNode]]] = [
    [ct(evolve(out_tree, use_fwd=True), cross_entropy_ablation_log_probs)],
    [ct(evolve(out_tree, use_fwd=True), get_cross_entropy_integrated_gradients_log_probs())],
    [evolve(out_tree, use_fwd=False)],
]

run_split = jax.jit(
    losses_runner_outputs(
        model,
        params,
        config=LossesRunnerTreeConfig(get_loss_no_deriv=True),
        get_extra_tree=lambda runner, seq_len: extra_tree,
    )
)

# %%

out_removed = stack_tree(
    run_on_tokens(
        run_split,
        data_subset,
        batch_size,
        disable_progress=False,
    )
)
out_removed

# %%

loss_removed = out_removed["deriv"]["loss"].mean(axis=0)
loss_orig = out_removed["no_deriv"]["loss"].mean()

ablation_loss = loss_removed[0]
print("ablation:", ablation_loss)
print("integrated gradient", loss_removed[1])
print("normal deriv", loss_removed[2])
print("original loss", loss_orig)

# %%

"""
Now let's start excluding some heads and see how the head does.
"""

# %%

tok_to_layer_1_attn_tree = output_tree(model, OutputConf(Embeds(), endpoint=AttnLayer(1)))
tok_to_layer_1_attn_tree

# %%

def all_except_run(
    q_heads,
    k_heads,
    v_heads,
    exclude_tok_to_k=False,
    ablation_over_ig_attn: Optional[bool] = True,
    ablation_over_ig_log_probs: Optional[bool] = True,
    config: LossesRunnerTreeConfig = LossesRunnerTreeConfig(avg_loss=False),
):
    def get_trees(runner, seq_len) -> ModifierCollectionTreeNodeStack:
        q_tree: ModifierCollectionTreeNodeStack = ModifierCollectionTreeNode(
            MulConf(ItemConf(KeyIdxs.single("blocks.attention.out_by_head", 0), ItemIdx(I[:, q_heads])))
        )
        k_tree: List[Optional[ModifierCollectionTreeNode]] = [
            ModifierCollectionTreeNode(
                MulConf(ItemConf(KeyIdxs.single("blocks.attention.out_by_head", 0), ItemIdx(I[:, k_heads])))
            ),
        ]
        if exclude_tok_to_k:
            k_tree.append(tok_to_layer_1_attn_tree)

        v_tree: ModifierCollectionTreeNodeStack = ModifierCollectionTreeNode(
            MulConf(ItemConf(KeyIdxs.single("blocks.attention.out_by_head", 0), ItemIdx(I[:, v_heads])))
        )

        base_tree = ct(
            ModifierCollectionTreeNode(
                StopGradConf(
                    ItemConf(KeyIdxs.single("blocks.attention.out_by_head", 1), ItemIdx(I[:, 4], except_idx=True))
                )
            ),
            ModifierCollectionTreeNode(StopGradConf(ItemConf(KeyIdxs.single("blocks.inp_res", 1)))),
            op.map(
                ablation_over_ig_attn,
                lambda x: ablation_softmax_probs(Idxs.single(1))
                if x
                else integrated_gradients_softmax_probs(Idxs.single(1)),
            ),
        )

        attn_mul: List[Optional[ModifierCollectionTreeNode]] = list(
            attn_score_mul_builder(
                runner,
                model,
                ct(base_tree, q_tree),
                ct(base_tree, k_tree),
                1,
            ).remove_conjunctive_from_sum()
        )

        all_mulled: List[Optional[ModifierCollectionTreeNode]] = list(
            probs_v_mul_builder(
                runner,
                attn_mul,
                ct(base_tree, v_tree),
                1,
            ).remove_conjunctive_from_sum()
        )

        return ct(
            base_tree,
            all_mulled,
            op.map(
                ablation_over_ig_log_probs,
                lambda x: cross_entropy_ablation_log_probs if x else get_cross_entropy_integrated_gradients_log_probs(),
            ),
        )

    return jax.jit(
        losses_runner_tree(
            partial_gpt_call(model, params, config=Gpt.CallConfig(log_finish=False)),
            get_trees,
            config=evolve(config, use_fwd=True),
        )
    )

# %%

# selected based on some experimentation
v_heads = list(range(1, 8))
base_q_heads = list(range(8))
base_k_heads = list(range(1, 8))

# %%

"""
Let's see how this does.
"""

# %%

run = all_except_run(base_q_heads, base_k_heads, v_heads)

# %%

out_except = stack_tree(run_on_tokens(run, data_subset, 2, disable_progress=False), stack=False)
mean_base = out_except["deriv"]["loss"].squeeze().mean(axis=(0, 1))

print(mean_base)
print("score", (ablation_loss - mean_base) / ablation_loss)

# %%

"""
Now let's try looping over including a single head on both k and q.
"""

# %%

for head_to_include in range(8):
    new_q_heads = copy(base_q_heads)
    new_q_heads.remove(head_to_include)
    new_k_heads = copy(base_k_heads)
    if head_to_include in new_k_heads:
        new_k_heads.remove(head_to_include)

    run = all_except_run(new_q_heads, new_k_heads, v_heads)

    # compile times on this are brutal...
    out_except = stack_tree(run_on_tokens(run, data[:100], 1, disable_progress=False), stack=False)
    out_except

    print("included:", head_to_include)
    mean_l = out_except["deriv"]["loss"].squeeze().mean(axis=(0, 1))
    print(mean_l)
    print("score", (ablation_loss - mean_l) / ablation_loss)

# %%

"""
Note that above we included layer 0 head 0 in v. It's somewhat suprising that
this head makes a reasonable contribution to v (but I claim it does).

Let's figure out exactly what this contribution is.

(This part of the notebook is undocumented and unexplained atm, but will
hopefully be better in the future after I write things up).
"""

# %%

def effect_of_layer_0_0_v(
    ablation_over_ig_log_probs: Optional[bool] = True,
    shape_on_toks=True,
    avg_loss=False,
    log_probs=True,
    fwd=True,
    toks_subset=I[:],
    extra_tree: ModifierCollectionTreeNodeStack = None,
):
    @jax.jit
    def f(toks, targets):
        runner = losses_runner_log(
            partial_gpt_call(model, params, config=Gpt.CallConfig(log_finish=False)), toks, targets
        )

        head_layer_0 = 0
        head_layer_1 = 4

        if not isinstance(toks_subset, slice):
            assert toks.shape[0] == 1

        shape = (
            (toks.shape[0], toks.shape[1] if isinstance(toks_subset, slice) else toks_subset.shape[0], 1)
            if shape_on_toks
            else ()
        )
        tree = ct(
            ModifierCollectionTreeNode(
                MulConf(ItemConf(KeyIdxs.single("blocks.attention.inp", 0), ItemIdx(I[:, toks_subset])), shape=shape)
            ),
            ModifierCollectionTreeNode(
                StopGradConf(
                    ItemConf(
                        KeyIdxs("blocks.attention.out_by_head", Idxs.single(0)),
                        item_idx=ItemIdx(I[:, head_layer_0], except_idx=True),
                    ),
                )
            ),
            ModifierCollectionTreeNode(StopGradConf(ItemConf(KeyIdxs.single("blocks.inp_res", 1)))),
            *stop_qkv_except(Idxs.single(1), "v"),
            ModifierCollectionTreeNode(
                StopGradConf(
                    ItemConf(
                        KeyIdxs.single("blocks.attention.out_by_head", 1),
                        ItemIdx(I[:, head_layer_1], except_idx=True),
                    ),
                )
            ),
            op.map(
                ablation_over_ig_log_probs,
                lambda x: cross_entropy_ablation_log_probs if x else get_cross_entropy_integrated_gradients_log_probs(),
            ),
            extra_tree,
        )

        if avg_loss:
            targets = [TargetConf(KeyIdxs("loss.mean_loss"), "loss")]
        else:
            targets = [TargetConf(KeyIdxs("loss.losses"), "loss")]

        if log_probs:
            targets.append(TargetConf(KeyIdxs("cross_entropy.log_probs"), display_name="log_probs"))

        return run_query(runner, Query(targets, tree, use_fwd=fwd))

    return f

# %%

run_probs_just_v = effect_of_layer_0_0_v(ablation_over_ig_log_probs=None)

# %%

out_except = stack_tree(
    run_on_tokens(run_probs_just_v, example_seq, 1, disable_progress=True, prepend_begin=False), stack=False
)
loss_diff = out_except["loss"].squeeze((1, 3))  # get 1 beyond to remove begin tok
summed_over = loss_diff.sum(-1)
_, idxs = jax.lax.top_k(jnp.abs(summed_over), k=10)
loss_diff[idxs]
np.array(example_seq_names)[np.array(idxs)], summed_over[idxs], summed_over

# %%

print(example_seq_names[97])
print_max_min_by_tok_k(loss_diff[97], k=20, get_tok=lambda idxs: np.array(example_seq_names)[np.array(idxs)])

# %%

print(example_seq_names[85])
print_max_min_by_tok_k(loss_diff[85], k=20, get_tok=lambda idxs: np.array(example_seq_names)[np.array(idxs)])

# %%

b_model = model.bind(params)

tok_embeds = b_model.embedding.token_embedding.embedding

all_ov_combined = b_model.get_ov_combined_mats_all_layers()
ov_0_0_to_1_4 = jnp.einsum(" o m, m i -> o i", all_ov_combined[1, 4], all_ov_combined[0, 0])
ov_0_0_to_1_4_to_out = jnp.einsum("v o, o i -> v i", tok_embeds, ov_0_0_to_1_4)
ov_0_0_to_out = jnp.einsum("v m, m i -> v i", tok_embeds, all_ov_combined[0, 0])
ov_0_0_to_1_4_to_out.shape, ov_0_0_to_out.shape

# %%

def print_for_tok(t: Union[str, int], k=50):
    if isinstance(t, str):
        t = single_tokenize(t)

    print(f'"{tok.decode(t)}"')

    preds_full = jnp.einsum("v i, i -> v", ov_0_0_to_1_4_to_out, tok_embeds[t])
    preds_0_0 = jnp.einsum("v i, i -> v", ov_0_0_to_out, tok_embeds[t])

    print("0_0_to_1_4")
    print("logit variance", preds_full.var())
    print_max_min_by_tok_k(preds_full, k=k)
    print()

    # print("0_0")
    # print("logit variance", preds_0_0.var())
    # print_max_min_by_tok_k(preds_0_0, k=k)

# %%

# print more toks as desired
print_for_tok(" dog", k=10)

# %%

# I think this isn't very meaningful due to general log softmax issues when averaging loss?
# Maybe not though.
run_probs_just_v_mean_loss = effect_of_layer_0_0_v(
    ablation_over_ig_log_probs=None,
    avg_loss=True,
    log_probs=False,
    fwd=False,
)

# %%

out_v_path = stack_tree(run_on_tokens(run_probs_just_v_mean_loss, data_subset, 2, disable_progress=False), stack=False)
losses_by_from_tok = out_v_path["loss"].squeeze(-1)
losses_by_from_tok.shape

# %%

n_tokens = tok_embeds.shape[0]

# %%

# offset by begin!
data_subset[:-1].shape, losses_by_from_tok[1:].shape

# %%

loss_deriv_by_tok = jnp.zeros(n_tokens).at[data_subset[:-1]].add(losses_by_from_tok[1:], unique_indices=False)

# %%

print_max_min_by_tok_k(loss_deriv_by_tok, k=30, normalize=False)

# %%

loss_deriv_by_tok.mean()

# %%

run_probs_from_current_loss = effect_of_layer_0_0_v(
    ablation_over_ig_log_probs=True, shape_on_toks=False, avg_loss=False, log_probs=False, fwd=True
)

# %%

batch_size = 2

# %%

out_v_path_q = stack_tree(
    run_on_tokens(run_probs_from_current_loss, data_subset, batch_size, disable_progress=False),
    stack=True,
)
losses_by_q_tok = einops.rearrange(out_v_path_q["loss"], "n (b s) -> (n b) s", b=batch_size)
losses_by_q_tok.shape

# %%

net_loss_by_q_tok = jnp.zeros(n_tokens).at[data_subset[:-1]].add(losses_by_q_tok[1:], unique_indices=False)
net_loss_by_q_tok.mean()

# %%

print_max_min_by_tok_k(net_loss_by_q_tok, k=30, normalize=False)

# %%

most_helpful = losses_by_q_tok.mean(axis=1).argmin()
most_helpful

# %%

print_max_min_by_tok_k(
    losses_by_q_tok[most_helpful][1:],
    k=20,
    get_tok=lambda idxs: toks_to_string_list(data_subset[most_helpful][(idxs,)]),
)

# %%

g_idxs = jnp.array((np.array(toks_to_string_list(data_subset[most_helpful])) == "g").nonzero()[0])
g_idxs

# %%

print(tok.decode(data_subset[most_helpful]))

# %%

run_probs_ablate_a_few_toks = effect_of_layer_0_0_v(
    ablation_over_ig_log_probs=True,
    shape_on_toks=True,
    avg_loss=False,
    log_probs=True,
    fwd=True,
    toks_subset=g_idxs + 2,  # + 1 for [BEGIN] and + 1 again for next tok
)

# %%

out_v_path_from_selected_toks = stack_tree(
    run_on_tokens(
        run_probs_ablate_a_few_toks,
        jnp.expand_dims(data_subset[most_helpful], 0),
        1,
        disable_progress=False,
    ),
    stack=True,
)
losses = out_v_path_from_selected_toks["loss"].squeeze((0, 2, -1))

# %%

losses[419]

# %%

all_virtual_head_outputs = jnp.einsum("o i, n i -> n o", ov_0_0_to_1_4, tok_embeds)
all_virtual_head_outputs.shape

# %%

import torch

x = torch.tensor(np.array(all_virtual_head_outputs))
u, s, v = torch.pca_lowrank(x, q=32)
u, s, v = u.numpy(), s.numpy(), v.numpy()
u.shape, s.shape, v.shape

# %%

# all_virtual_head_outputs[0]
idx = 21
jnp.abs(all_virtual_head_outputs[idx] - (v @ jnp.transpose(v)) @ all_virtual_head_outputs[idx]).max()

# %%

# inp_to_dims
inp_to_dims = jnp.einsum("o i, o d -> d i", ov_0_0_to_1_4, v)
tok_input_to_dims_values = jnp.einsum("d i, t i -> d t", inp_to_dims, tok_embeds)
inp_to_dims.shape, tok_input_to_dims_values.shape

# %%

plt.plot(s)

# %%

s

# %%

direc_to_toks = jnp.einsum("v o, o d -> d v", tok_embeds, v)

# %%

for dim in range(5):
    print("dim", dim)
    print_max_min_by_tok_k(direc_to_toks[dim], k=20)

# %%

print(print_max_min_by_tok_k(tok_input_to_dims_values[0], k=20))

# %%

def get_run_probs_ablate_some_dims(up_to_n):
    return effect_of_layer_0_0_v(
        ablation_over_ig_log_probs=True,
        shape_on_toks=False,
        avg_loss=True,
        log_probs=False,
        fwd=True,
        extra_tree=ModifierCollectionTreeNode(
            ModifyInBasisConf(
                ItemConf(KeyIdxs.single("blocks.attention.out_by_head", 1), ItemIdx(I[:, 4])),
                proj_to=jnp.transpose(v),
                proj_back=v,
                sub_conf=StopGradConf(ItemConf(KeyIdxs(""), ItemIdx(I[..., :up_to_n]))),
            )
        ),
    )

# %%

run_ablate_after_1 = get_run_probs_ablate_some_dims(1)

# %%

out_some_dims = stack_tree(run_on_tokens(run_ablate_after_1, data_subset, 2, disable_progress=True), stack=True)
out_some_dims["loss"], out_some_dims["loss"].shape, out_some_dims["loss"].mean()

# %%

losses = [
    stack_tree(run_on_tokens(get_run_probs_ablate_some_dims(i), data_subset, 2, disable_progress=False), stack=True)[
        "loss"
    ].mean()
    for i in range(0, 10)
]

# %%

plt.plot(-np.array(losses))


import pickle
from gensim.models import Word2Vec

with open("data/word_vector.pkl", "rb") as f:
    word_vector = pickle.load(f)

#train word2vec model
model = Word2Vec(sentences=word_vector, window=3, sg=1, hs=0, negative=20,
                 alpha=0.03, min_alpha=0.0007, vector_size=100, epochs=10,
                 seed=24)

#save the model
with open("model/word2vec_model.pkl", "wb") as f:
    pickle.dump(model, f, protocol=pickle.HIGHEST_PROTOCOL)


from django.db import models
from django.contrib.auth.models import User
from django.utils.translation import ugettext_lazy as _
from django.core.validators import MaxValueValidator, MinValueValidator

# Create your models here.

CATEGORY = (
        ('Education', 'Education'),
        ('Healthcare', 'Healthcare'),
        ('Hospitality', 'Hospitality'),
        ('Manufacturing', 'Manufacturing'),
        ('Media', 'Media'),
        ('Software', 'Software'),
    )

class Company(models.Model):
	admin = models.OneToOneField(User, on_delete=models.SET_NULL, null=True)
	name = models.CharField(_("Company Name"), max_length=200, unique=True)
	category = models.CharField(choices=CATEGORY, max_length=100)
	email = models.CharField(max_length=100, unique=True)
	mobile = models.PositiveIntegerField(blank=True, null=True, validators=[MaxValueValidator(9999999999)])
	address = models.TextField(max_length=200, blank=True, null=True)

	def __str__(self):
		return self.name

class UserProfile(models.Model):
	user = models.OneToOneField(User, on_delete=models.CASCADE)
	name = models.CharField(max_length=255)
	date_of_birth = models.DateField(_("Date of birth"), blank=True, null=True)
	email = models.EmailField(max_length=255, unique=True, null=True)
	mobile = models.PositiveIntegerField(blank=True, null=True, validators=[MaxValueValidator(9999999999)])
	company = models.ForeignKey(Company, on_delete=models.CASCADE)

	def __str__(self):
		return self.user.username


import code
import readline
import atexit
import os

class HistoryConsole(code.InteractiveConsole):
    def __init__(self, locals=None, filename="<console>",
                 histfile=os.path.expanduser("~/.console-history")):
        code.InteractiveConsole.__init__(self, locals, filename)
        self.init_history(histfile)

    def init_history(self, histfile):
        readline.parse_and_bind("tab: complete")
        if hasattr(readline, "read_history_file"):
            try:
                readline.read_history_file(histfile)
            except IOError:
                pass
            atexit.register(self.save_history, histfile)

    def save_history(self, histfile):
        readline.set_history_length(1000)
        readline.write_history_file(histfile)

if __name__ == '__main__':
    while True:
        c = HistoryConsole().raw_input()
        print(c)


#!/usr/bin/env python3

import sys
import asyncio
import aiohttp
import json

namespaces = {}
broker = "tcp://localhost:1883"

###############################################################################
####################################################################### helpers

def valid_python_version ():
    v = sys.version_info
    if v[0] != 3: return False
    if v[1] < 5: return False
    return True

###############################################################################
#################################################################### processing

async def attach_mqtt_dao (prefix, entity, topic):
    topic = '/building1/'+topic

    MqttLiveData    = await model_ensure_instance(prefix, 'dao:MqttLiveData' , escape('data'+topic))
    hasBroker       = await model_ensure_relationship(MqttLiveData , 'dao:hasBroker', '"%s"' % broker)
    hasTopic        = await model_ensure_relationship(MqttLiveData , 'dao:hasTopic', '"%s"' % topic)
    hasMqttLiveData = await model_ensure_relationship(entity , 'dao:hasMqttLiveData', MqttLiveData)

async def process ():
    # instances
    prefix = 'n:'
    dh   = await model_ensure_instance(prefix, 'gfb:District_Heat' , escape('district_heat'))
    hx   = await model_ensure_instance(prefix, 'gfb:Heat_Exchanger', escape('heat_exchanger'))
    tank      = await model_ensure_instance(prefix, 'gfb:Hot_Water_Tank'          , escape('hot_water_tank'))
    tank_temp = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('hot_water_tank_temperature'))
    d_ret_flow  = await model_ensure_instance(prefix, 'gfb:Water_Flow_Sensor'                 , escape('district_flow'))
    d_ret_temp  = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor'          , escape('district_return_temp'))
    d_sup_temp  = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor'          , escape('district_supply_temp'))
    d_ret_valve = await model_ensure_instance(prefix, 'gfb:Valve'                             , escape('district_return_valve'))
    d_ret_pres  = await model_ensure_instance(prefix, 'gfb:Water_Pressure_Sensor'             , escape('district_return_pressure'))
    d_sup_pres  = await model_ensure_instance(prefix, 'gfb:Water_Pressure_Sensor'             , escape('district_supply_pressure'))
    d_dif_pres  = await model_ensure_instance(prefix, 'gfb:Water_Differential_Pressure_Sensor', escape('district_differential_pressure'))
    d_heat      = await model_ensure_instance(prefix, 'gfb:Heat_Meter'                        , escape('district_heat_meter'))
    h_sup_temp = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('heated_supply_temp'))
    h_l1_sup_valve = await model_ensure_instance(prefix, 'gfb:Valve'                   , escape('heated_loop1_supply_valve'))
    h_l1_sup_temp  = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('heated_loop1_supply_temp'))
    h_l1_sup_pump  = await model_ensure_instance(prefix, 'gfb:Pump'                    , escape('heated_loop1_supply_pump'))
    h_l1_radiator  = await model_ensure_instance(prefix, 'gfb:Radiator'                , escape('heated_loop1_radiator'))
    h_l2_sup_bvalve    = await model_ensure_instance(prefix, 'gfb:Bypass_Valve'            , escape('heated_loop2_supply_bypass_valve'))
    h_l2_sup_temp      = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('heated_loop2_supply_temp'))
    h_l2_ret_pump      = await model_ensure_instance(prefix, 'gfb:Pump'                    , escape('heated_loop2_return_pump'))
    h_l2_ret_temp_pre  = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('heated_loop2_return_prebypass_temp'))
    h_l2_ret_temp_post = await model_ensure_instance(prefix, 'gfb:Water_Temperature_Sensor', escape('heated_loop2_return_postbypass_temp'))
    h_l2_radiator      = await model_ensure_instance(prefix, 'gfb:Radiator'                , escape('heated_loop2_radiator'))
    h_l3_ret_valve = await model_ensure_instance(prefix, 'gfb:Valve', escape('heated_loop3_return_valve'))

    # attach data sources
    await attach_mqtt_dao(prefix, tank_temp  , 'tank/temp')
    await attach_mqtt_dao(prefix, d_ret_flow , 'district/return/flow')
    await attach_mqtt_dao(prefix, d_ret_temp , 'district/return/temp')
    await attach_mqtt_dao(prefix, d_sup_temp , 'district/supply/temp')
    await attach_mqtt_dao(prefix, d_ret_valve, 'district/return/valve/position')
    await attach_mqtt_dao(prefix, d_ret_pres , 'district/return/pressure')
    await attach_mqtt_dao(prefix, d_sup_pres , 'district/supply/pressure')
    await attach_mqtt_dao(prefix, d_dif_pres , 'district/differential/pressure')
    await attach_mqtt_dao(prefix, d_heat     , 'district/heat')
    await attach_mqtt_dao(prefix, h_sup_temp, 'heated/supply/temp')
    await attach_mqtt_dao(prefix, h_l1_sup_valve, 'heated/loop1/supply/valve/position') # TODO: should be attached to a sensor on the valve
    await attach_mqtt_dao(prefix, h_l1_sup_temp , 'heated/loop1/supply/temp')
    await attach_mqtt_dao(prefix, h_l1_sup_pump , 'heated/loop1/supply/pump/rpm') # TODO: should be attached to a sensor on the pump
    await attach_mqtt_dao(prefix, h_l2_sup_bvalve   , 'heated/loop2/supply/bypass-valve/position') # TODO: should be attached to a sensor on the bypass valve
    await attach_mqtt_dao(prefix, h_l2_sup_temp     , 'heated/loop2/supply/temp')
    await attach_mqtt_dao(prefix, h_l2_ret_pump     , 'heated/loop2/return/pump/rpm') # TODO: should be attached to a sensor on the pump
    await attach_mqtt_dao(prefix, h_l2_ret_temp_pre , 'heated/loop2/return/temp/pre')
    await attach_mqtt_dao(prefix, h_l2_ret_temp_post, 'heated/loop2/return/temp/post')
    await attach_mqtt_dao(prefix, h_l3_ret_valve, 'heated/loop3/return/valve/position') # TODO: should be attached to a sensor on the valve

    # relations: district heated water
    await model_ensure_relationship(dh         , 'gfb:feedsSupplyDistrictHeatedWater', d_sup_temp)
    await model_ensure_relationship(d_sup_temp , 'gfb:feedsSupplyDistrictHeatedWater', d_sup_pres)
    await model_ensure_relationship(d_sup_pres , 'gfb:feedsSupplyDistrictHeatedWater', hx)
    await model_ensure_relationship(hx         , 'gfb:feedsReturnDistrictHeatedWater', d_ret_pres)
    await model_ensure_relationship(d_ret_pres , 'gfb:feedsReturnDistrictHeatedWater', d_ret_valve)
    await model_ensure_relationship(d_ret_valve, 'gfb:feedsReturnDistrictHeatedWater', d_ret_temp)
    await model_ensure_relationship(d_ret_temp , 'gfb:feedsReturnDistrictHeatedWater', d_ret_flow)
    await model_ensure_relationship(d_sup_temp, 'gfb:controls', d_heat)
    await model_ensure_relationship(d_ret_temp, 'gfb:controls', d_heat)
    await model_ensure_relationship(d_ret_flow, 'gfb:controls', d_heat)
    await model_ensure_relationship(d_sup_pres, 'gfb:controls', d_dif_pres)
    await model_ensure_relationship(d_ret_pres, 'gfb:controls', d_dif_pres)

    # relations: heated
    await model_ensure_relationship(hx, 'gfb:feedsSupplyHeatedWater', h_sup_temp)

    # relations: heated loop 1
    await model_ensure_relationship(h_sup_temp    , 'gfb:feedsSupplyHeatedWater', h_l1_sup_valve)
    await model_ensure_relationship(h_l1_sup_valve, 'gfb:feedsSupplyHeatedWater', h_l1_sup_temp)
    await model_ensure_relationship(h_l1_sup_temp , 'gfb:feedsSupplyHeatedWater', h_l1_sup_pump)
    await model_ensure_relationship(h_l1_sup_pump , 'gfb:feedsSupplyHeatedWater', h_l1_radiator)
    await model_ensure_relationship(h_l1_radiator , 'gfb:feedsReturnHeatedWater', hx)

    # relations: heated loop 2
    await model_ensure_relationship(h_sup_temp        , 'gfb:feedsSupplyHeatedWater', h_l2_sup_bvalve)
    await model_ensure_relationship(h_l2_sup_bvalve   , 'gfb:feedsSupplyHeatedWater', h_l2_sup_temp)
    await model_ensure_relationship(h_l2_sup_temp     , 'gfb:feedsSupplyHeatedWater', h_l2_radiator)
    await model_ensure_relationship(h_l2_radiator     , 'gfb:feedsReturnHeatedWater', h_l2_ret_pump)
    await model_ensure_relationship(h_l2_ret_pump     , 'gfb:feedsReturnHeatedWater', h_l2_ret_temp_pre)
    await model_ensure_relationship(h_l2_ret_temp_pre , 'gfb:feedsReturnHeatedWater', h_l2_sup_bvalve)
    await model_ensure_relationship(h_l2_ret_temp_pre , 'gfb:feedsReturnHeatedWater', h_l2_ret_temp_post)
    await model_ensure_relationship(h_l2_ret_temp_post, 'gfb:feedsReturnHeatedWater', hx)

    # relations: heated loop 3
    await model_ensure_relationship(h_sup_temp    , 'gfb:feedsSupplyHeatedWater', tank)
    await model_ensure_relationship(tank_temp     , 'brick:isPointOf'           , tank)
    await model_ensure_relationship(tank          , 'gfb:feedsReturnHeatedWater', h_l3_ret_valve)
    await model_ensure_relationship(h_l3_ret_valve, 'gfb:feedsReturnHeatedWater', hx)

    # store resulting model
    await rdf_store()

def escape (value: str):
    value = value.replace(' ', '_space_')
    value = value.replace(':', '_colon_')
    value = value.replace('/', '_slash_')
    value = value.replace('æ', '_ae_')
    value = value.replace('ø', '_ao_')
    value = value.replace('å', '_aa_')
    value = value.replace('Æ', '_AE_')
    value = value.replace('Ø', '_AO_')
    value = value.replace('Å', '_AA_')
    return value

###############################################################################
######################################################################### model

async def model_ensure_instance (prefix: str, datatype: str, label: str):
    q = '''
    SELECT ?name
    WHERE {
        ?name rdf:type %s .
        ?name rdf:label "%s" .
    }
    ''' % (datatype, label)
    success, rs = await rdf_query(q)

    if success and len(rs)>0:
        name = rs[0][0]
        if '#' in name:
            parts = name.split('#')
            if len(parts)!=2:
                print('ERROR: Unknown format of entity "%s".' % name)
            parts[0] += '#'
            if parts[0] in namespaces:
                name = '%s:%s' % (namespaces[parts[0]], parts[1])
            else:
                print('ERROR: Namespace not defined for entity "%s".' % name)
    else:
        name = '%s_%s' % (prefix, label)
        await rdf_update_split(insert_clause=['%s rdf:type %s' % (name, datatype),
                                              '%s rdf:label "%s"' % (name, label)])
    return name

async def model_ensure_relationship (sub: str, pred: str, obj: str, onlyobj=False):
    if onlyobj:
        await rdf_update_split(delete_clause=['%s %s ?obj' % (sub, pred)],
                               where_clause=['%s %s ?obj' % (sub, pred)])

    await rdf_update_split(insert_clause=['%s %s %s' % (sub, pred, obj)])

###############################################################################
########################################################################### rdf

async def rdf_namespaces ():
    url = 'http://%s:%u/namespaces' % (host, port)
    async with aiohttp.ClientSession() as session:
        async with session.post(url, data='"dummy"') as response:
            text = await response.text()

            if response.status != 200:
                print('ERROR: Unable to lookup namespaces')
                print(text)
                return False, {}

            try:
                return True, json.loads(text)
            except Exception as e:
                print('ERROR: Exception while trying to parse result of namespace lookup')
                print(text)
                print(str(e))
                return False, {}

async def rdf_store ():
    url = 'http://%s:%u/store' % (host, port)
    async with aiohttp.ClientSession() as session:
        async with session.post(url, data='"dummy"') as response:
            text = await response.text()

            if response.status != 200:
                print('ERROR: Unable to store model')
                print(text)
                return False

            print('NOTICE: Successfully stored model')
            return True

async def rdf_query (query: str):
    query = json.dumps(query, sort_keys=True, indent=4, separators=(',', ': '))
    url = 'http://%s:%u/query' % (host, port)
    async with aiohttp.ClientSession() as session:
        async with session.put(url, data=query) as response:
            text = await response.text()

            if response.status != 200:
                print('ERROR: Unable to query model')
                print(query)
                print(text)
                1/0
                return False, None

            try:
                return True, json.loads(text)['resultset']
            except Exception as e:
                print('ERROR: Exception while trying to parse result model query')
                print(text)
                print(str(e))
                return False, None

async def rdf_update (query: str):
    query = json.dumps(query, sort_keys=True, indent=4, separators=(',', ': '))
    url = 'http://%s:%u/update' % (host, port)
    async with aiohttp.ClientSession() as session:
        async with session.post(url, data=query) as response:
            text = await response.text()

            if response.status != 200:
                print('ERROR: Unable to update model with status "%d"' % response.status)
                print(text)
                1/0
                return False

            try:
                r = json.loads(text)
                if not 'success' in r or type(r['success'])!=bool:
                    print('ERROR: Unable to parse result of model update')
                    print(text)
                    return False

                print('NOTICE: Success of model update: '+str(r['success']))
                return r['success']
            except:
                print('ERROR: Exception while trying to parse result of model update')
                print(text)
                print(str(e))
                return False

async def rdf_update_split (ns : list = None, insert_clause : str = None, delete_clause : str = None, where_clause : str = None):
    if ns == None: ns = []

    q = []

    for entry in ns:
        q.append('%s' % entry)
    q.append('')

    if insert_clause:
        q.append('INSERT {')
        for entry in insert_clause:
            q.append('    %s .' % entry)
        q.append('}')

    if delete_clause:
        q.append('DELETE {')
        for entry in delete_clause:
            q.append('    %s .' % entry)
        q.append('}')

    q.append('WHERE {')
    if where_clause:
        for entry in where_clause:
            q.append('    %s .' % entry)
    q.append('}')

    q = ''.join(map(lambda line: '%s\n' % line, q))
    print(q)

    r = await rdf_update(q)
    print('response:', r)

###############################################################################
########################################################################## main

async def main():
    global namespaces

    # load namespaces
    print('STATUS: Loading namespaces:')
    success, ns = await rdf_namespaces()
    if not success or not 'success' in ns or not ns['success']:
        print('ERROR: Unable to fetch namespaces:')
        print(success)
        print(ns)
        sys.exit(3)
    for key in ns['namespaces']:
        prefix = ns['namespaces'][key]
        print('STATUS: - %s : %s' % (prefix, key))
        namespaces[prefix] = key

    await process()

# guard: python version
if not valid_python_version():
    print('ERROR: Invalid python version (%s), bust be 3.(5+).' % str(sys.version_info))
    sys.exit(1)

# guard: commandline arguments
if len(sys.argv) != 4:
    print('Syntax: %s NAMESPACE RDF_SERVER_HOST RDF_SERVER_PORT' % sys.argv[0])
    print('        %s http://ss.sdu.dk/test/grundstrup-bidi/20200727/# 127.0.0.1 8001' % sys.argv[0])
    sys.exit(2)

# extract parameters
namespace   =     sys.argv[1]
host        =     sys.argv[2]
port        = int(sys.argv[3])

loop = asyncio.get_event_loop()

# enter service loop
try:
    loop.run_until_complete(main())
except KeyboardInterrupt:
    print('')
    print('STATUS: Exiting ...')
    loop.close()
    exit(0)

########################################################################### EOF
###############################################################################


"""Friend request model."""

# Django
from django.db import models

# Utilities
from app.utils.models import FbModel

class FriendRequest(FbModel):
    """Friend request."""

    requesting_user = models.ForeignKey('users.User', on_delete=models.CASCADE)

    requested_user = models.ForeignKey(
        'users.User',
        on_delete=models.CASCADE,
        related_name='requested_user')

    accepted = models.BooleanField(default=False)

    def __str__(self):
        """Return requesting_user and requested_user."""
        return 'from @{} to @{}'.format(
            self.requesting_user.username,
            self.requested_user.username)


#!/usr/bin/env python

"""
Common utilities that can be widely used across the package
"""

import numpy as np
import os
import glob
from itertools import groupby
import sys
import time

from icecube.utils.logger import Logger

logger = Logger(os.path.basename(__file__))

class DirUtils:
    @staticmethod
    def get_dir_files(local_dir, fext=None):
        """
        scrape the given directory for all files or files with certain extension.
        fext: str or list of strings.
        """
        local_dir = os.path.join(local_dir, "")
        fpaths = []  # contains names of all images.

        if fext and not (isinstance(fext, list)):
            fext = list(fext)

        if fext:
            for ext in fext:
                ext = ext.replace("*", "")  # if asterik was passed by mistake.
                fpaths.extend(glob.glob(local_dir + "*" + ext))
        else:
            fpaths.extend(glob.glob(local_dir + "*"))

        fnames = [p.replace(local_dir, "") for p in fpaths]

        return fnames, fpaths

def get_dummy_metadata():
    """
    Incase of a temporal gap, the following metadata is used
    """
    return dict({"product_file": "None"})

def get_product_metadata(product_file):
    """
    Metadata appended for labels xr.DataArray
    """
    return dict({"product_file": product_file})

def assert_metadata_exists(metadata_df):
    """
    Make sure that metadata is non-empty
    Returns True if non-empty otherwise raises exception
    """
    for df_col in metadata_df.columns:
        if metadata_df[df_col].isnull().values.all():
            raise ValueError(
                "No valid values found for pd.Series: {} for pd.column:{} against given configuration".format(
                    metadata_df[df_col], df_col
                )
            )

    return True

def all_equal(iterable):
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

def match_dict_key(d, value):
    """
    get key by matching value of a dictionary.
    d: dictionary of type dict.
    value: value to match
    returns: key of matched value.
    """
    for k, v in d.items():
        if v == value:
            return k

    print(f"value: {value} not found in dictionary")

def draw_progress_bar(percent, barLen=50):
    sys.stdout.write("\r")
    progress = ""
    for i in range(barLen):
        if i < int(barLen * percent):
            progress += "="
        else:
            progress += " "
    sys.stdout.write(f"[ {progress} ] {percent * 100:.2f}%")
    sys.stdout.flush()

def measure_time(func):
    """
    :param func:
    :return:
    """

    def time_it(*args, **kwargs):
        time_started = time.time()
        return_value = func(*args, **kwargs)
        time_elapsed = time.time()
        logger.info(
            "{execute} running time is {sec} seconds".format(
                execute=func.__name__,
                sec=round(time_elapsed - time_started, 4),
            )
        )
        return return_value

    return time_it

class NumpyEncoder:
    """Special json encoder for numpy types"""

    @classmethod
    def encode(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return str(obj.tolist())
        else:
            return obj

def get_slc_metadata(h5f):
    """
    Read metadata from SLC header

    Args:
        opened_h5_file: SLC h5 file opened as a h5py File
    """

    meta_dict = {}

    # h5py doesn't read .keys() correctly, but returns a generator. So parsing to list to get all keys.
    key_list = [x for x in h5f.keys()]

    # Check if there already is a list of bands
    if "bands" in key_list:
        non_meta_keys = h5f["bands"]
    else:
        non_meta_keys = ["s_i", "s_q"]

    # Splines are not passed as they are not necessary. We just skip them for now.
    # This doesn't not effect the SLC in any meaningfull way
    additional_bands_to_skip_for_now = [
        "RPC",
        "height_spline",
        "lat_spline",
        "lon_spline",
    ]
    for key in key_list:
        if key not in non_meta_keys + additional_bands_to_skip_for_now:

            h5_meta_val = h5f[key]

            if type(h5_meta_val[()]) == bytes:
                # h5py 3.0+ introduces str stored as bytes. We use .asstr() to read it (but breaks with h5py < 2.10)
                meta_dict[key] = np.array_str(np.array(h5_meta_val.asstr()[()]))

            elif type(h5_meta_val[()]) == np.bytes_:
                # This is most likely a date
                meta_dict[key] = np.array_str(
                    np.squeeze(np.char.decode(h5_meta_val, "utf-8"))
                )

            else:
                # We're handling a numpy array so we should be good to go
                meta_dict[key] = np.array_str(np.array(h5_meta_val[()]))

                # RPCs are nested under "RPC/" in the h5 thus need to be parsed in a specific manner
    RPC_source = h5f["RPC"]

    for key, val in RPC_source.items():
        meta_dict[f"RPC_{key}"] = np.array(val, dtype=np.float32)

    return meta_dict


import ast
import os
import re
import sys

from setuptools import setup, __version__ as setuptools_version

def readme(name='README.rst'):
    try:
        with open(name) as f:
            rst = f.read()
        return re.sub(
            r'(^|\n).. include::\s*([^\n]+)($|\n)',
            lambda m: m.group(1) + (readme(m.group(2)) or '') + m.group(3),
            rst
        )
    except (IOError, OSError):
        return

def get_version():
    module_path = os.path.join(os.path.dirname(__file__), 'nirum_http.py')
    module_file = open(module_path)
    try:
        module_code = module_file.read()
    finally:
        module_file.close()
    tree = ast.parse(module_code, module_path)
    for node in ast.iter_child_nodes(tree):
        if not isinstance(node, ast.Assign) or len(node.targets) != 1:
            continue
        target, = node.targets
        if isinstance(target, ast.Name) and target.id == '__version__':
            value = node.value
            if isinstance(value, ast.Str):
                return value.s
            raise ValueError('__version__ is not defined as a string literal')
    raise ValueError('could not find __version__')

setup_requires = []
install_requires = [
    'nirum >= 0.6.0',
    'requests >= 2.11.1',
    'six',
]
tests_require = [
    'flake8-import-order >= 0.12, < 1.0',
    'flake8-import-order-spoqa >= 1.0.0, < 2.0.0',
    # flake8 does not yet support pycodestyle 2.4.0.
    # Can be remove after this issue is fixed: https://gitlab.com/pycqa/flake8/issues/415
    'pycodestyle >= 2.0, < 2.4.0',
    'pytest >= 3.1.2, < 4.0.0',
    'pytest-flake8 >= 0.9.0, < 1.0.0',
    'requests-mock >= 1.3.0, < 1.4.0',
]
extras_require = {
    'tests': tests_require,
}
below35_requires = [
    'typing',
]

if 'bdist_wheel' not in sys.argv and sys.version_info < (3, 5):
    install_requires.extend(below35_requires)

if tuple(map(int, setuptools_version.split('.'))) < (17, 1):
    setup_requires = ['setuptools >= 17.1']
    extras_require.update({":python_version=='3.4'": below35_requires})
    extras_require.update({":python_version=='2.7'": below35_requires})
else:
    extras_require.update({":python_version<'3.5'": below35_requires})

setup(
    name='nirum-http',
    version=get_version(),
    description='Nirum HTTP transport for Python',
    long_description=readme(),
    url='https://github.com/nirum-lang/nirum-python-http',
    bugtrack_url='https://github.com/nirum-lang/nirum/issues',
    author='Kang Hyojun',
    author_email='iam.kanghyojun' '@' 'gmail.com',
    license='MIT license',
    py_modules=['nirum_http'],
    install_requires=install_requires,
    setup_requires=setup_requires,
    extras_require=extras_require
)


import torch
import torch.nn as nn
import numpy as np

class InferenceRunner:

    def __init__(self, model, prediction_mode='simple', scaler_x=None, scaler_y=None):
        self.model = model
        self.prediction_mode = prediction_mode
        self.last_prediction = None
        self.scaler_x = scaler_x
        self.scaler_y = scaler_y

    def predict_rec(self, features):
        pass

    def predict_simple(self, features):
        features = np.stack([self.scaler_x.transform(features.cpu()[i]) for i in range(features.shape[0])], 0).astype(np.float32)
        features = torch.from_numpy(features).to('cuda' if torch.cuda.is_available() else 'cpu')
        preds = self.model(features)
        if self.scaler_y:
            preds = self.scaler_y.inverse_transform(preds.detach().cpu().numpy())
        return preds

    def __call__(self, features):
        if self.prediction_mode == 'recurrent':
            return self.predict_rec(features)
        elif self.prediction_mode == 'simple':
            return self.predict_simple(features)
        else:
            raise ValueError('Invalid prediction mode. Expected recurrent or simple')

class LSTMPredictor(nn.Module):

    def __init__(self, num_features, hidden_dim, tagset_size=6):
        super(LSTMPredictor, self).__init__()
        self.num_layers = 1
        self.hidden_dim = hidden_dim
        self.lstm = nn.LSTM(num_features, hidden_dim, batch_first=True)
        self.sim_prediction = nn.Linear(hidden_dim, tagset_size)

    def forward(self, features):
        h0 = torch.zeros(1, features.size(0), self.hidden_dim).to('cuda' if torch.cuda.is_available() else 'cpu')
        c0 = torch.zeros(1, features.size(0), self.hidden_dim).to('cuda' if torch.cuda.is_available() else 'cpu')
        seq, (hidden_state, cell_state) = self.lstm(features, (h0, c0))
        hidden_state = hidden_state.detach()
        predictions = self.sim_prediction(hidden_state.squeeze())
        predictions = predictions + features[:, -1].squeeze()
        return predictions

class Conv1dPredictor(nn.Module):

    def __init__(self, num_features, target_size=6):
        super(Conv1dPredictor, self).__init__()
        self.conv1 = nn.Conv1d(num_features, 256, 2, stride=1)
        self.conv2 = nn.Conv1d(256, 256, 2, stride=1)
        self.conv3 = nn.Conv1d(256, 256, 2, stride=1)
        self.conv4 = nn.Conv1d(256, target_size, 2, stride=1)
        self.preds = nn.Linear(21 * target_size, target_size)

        self.bn1 = nn.BatchNorm1d(256)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(target_size)

        self.pool1 = nn.MaxPool1d(1)
        self.pool2 = nn.MaxPool1d(1)
        self.pool3 = nn.MaxPool1d(1)
        self.pool4 = nn.MaxPool1d(1)

    def forward(self, features):
        x = self.conv1(features.permute(0, 2, 1))
        x = self.bn1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.pool3(x)
        x = self.conv4(x)
        x = self.bn4(x)
        x = self.pool4(x)
        predictions = self.preds(x.view(x.size(0), -1))
        return predictions

class MLPPredictor(nn.Module):

    def __init__(self, num_features, seq_len, hidden_dim, target_size=6):
        super(MLPPredictor, self).__init__()
        self.num_layers = 1
        self.hidden_dim = hidden_dim
        self.num_features = num_features
        self.seq_len = seq_len
        self.dense1 = nn.Linear(seq_len * num_features, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.dense2 = nn.Linear(hidden_dim, hidden_dim // 2)
        self.bn2 = nn.BatchNorm1d(hidden_dim // 2)
        self.dense3 = nn.Linear(hidden_dim // 2, target_size)
        self.target_size = target_size
        self.relu = nn.ReLU()

    def forward(self, features):
        x = features.view(-1, self.seq_len * self.num_features)
        x = self.dense1(x)
        # x = self.bn1(x)
        x = self.dense2(x)
        # x = self.bn2(x)
        predictions = self.dense3(x)
        predictions = predictions.view(-1, self.target_size)
        return predictions