from raybot import config
from raybot.model import db, POI, Location
from raybot.bot import bot, dp
from raybot.util import h, HTML, split_tokens, get_buttons, get_map, get_user, tr, DOW
from raybot.actions.poi import POI_EDIT_CB, POI_LIST_CB
from raybot.actions.messages import broadcast_str, broadcast
import re
import os
import logging
import random
import humanized_opening_hours as hoh
from aiosqlite import DatabaseError
from string import ascii_lowercase
from datetime import datetime
from typing import Dict, Union
from collections import Counter
from aiogram import types
from aiogram.utils.callback_data import CallbackData
from aiogram.dispatcher import FSMContext
from aiogram.dispatcher.filters.state import State, StatesGroup
from aiogram.utils.exceptions import TelegramAPIError, MessageToDeleteNotFound

HOUSE_CB = CallbackData('ehouse', 'hid')
FLOOR_CB = CallbackData('efloor', 'floor')
BOOL_CB = CallbackData('boolattr', 'attr', 'value')
PHOTO_CB = CallbackData('ephoto', 'name', 'which')
TAG_CB = CallbackData('etag', 'tag')
TAG_PAGE_CB = CallbackData('tagpg', 'page')

class EditState(StatesGroup):
    name = State()
    location = State()
    keywords = State()
    confirm = State()
    attr = State()
    comment = State()
    message = State()

def cancel_keyboard():
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton('❌ ' + tr('cancel'), callback_data='cancel')
    )

def location_keyboard():
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton(
            tr(('editor', 'latlon')),
            url='https://zverik.github.io/latlon/#16/53.9312/27.6525'),
        types.InlineKeyboardButton('❌ ' + tr('cancel'), callback_data='cancel'),
    )

def valid_location(loc):
    bbox = config.BBOX
    if not bbox or len(bbox) != 4:
        return True
    return bbox[0] <= loc.lon <= bbox[2] and bbox[1] <= loc.lat <= bbox[3]

@dp.callback_query_handler(state=EditState.all_states, text='cancel')
async def new_cancel(query: types.CallbackQuery, state: FSMContext):
    await delete_msg(query, state)
    await state.finish()
    user = await get_user(query.from_user)
    if user.review:
        kbd = types.InlineKeyboardMarkup().add(types.InlineKeyboardButton(
            '🗒 ' + tr(('review', 'continue')), callback_data='continue_review'))
    else:
        kbd = get_buttons()
    await bot.send_message(
        query.from_user.id,
        tr(('new_poi', 'cancel')),
        reply_markup=kbd
    )

@dp.callback_query_handler(state='*', text='new')
async def new_poi(query: types.CallbackQuery):
    if config.MAINTENANCE:
        await bot.send_message(query.from_user.id, tr('maintenance'))
        return
    await EditState.name.set()
    await bot.send_message(
        query.from_user.id,
        tr(('new_poi', 'name')),
        reply_markup=cancel_keyboard()
    )

@dp.callback_query_handler(POI_EDIT_CB.filter(), state='*')
async def edit_poi(query: types.CallbackQuery, callback_data: Dict[str, str],
                   state: FSMContext):
    if config.MAINTENANCE:
        await bot.send_message(query.from_user.id, tr('maintenance'))
        return
    if callback_data['d'] == '1':
        await delete_msg(query)
    poi = await db.get_poi_by_id(int(callback_data['id']))
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(query.from_user, state)

@dp.message_handler(state=EditState.name)
async def new_name(message: types.Message, state: FSMContext):
    name = message.text.strip()
    if len(name) < 3:
        await message.answer(tr(('new_poi', 'name_too_short')))
        return
    await state.set_data({'name': name})
    await EditState.location.set()
    await message.answer(tr(('new_poi', 'location')), reply_markup=location_keyboard())

def parse_location(message: types.Message):
    ll = re.match(r'^\s*(-?\d+\.\d+),\s*(-?\d+\.\d+)\s*$', message.text or '')
    if message.location:
        return Location(lon=message.location.longitude, lat=message.location.latitude)
    elif ll:
        return Location(lat=float(ll.group(1)), lon=float(ll.group(2)))
    return None

@dp.message_handler(state=EditState.location,
                    content_types=[types.ContentType.TEXT, types.ContentType.LOCATION])
async def new_location(message: types.Message, state: FSMContext):
    loc = parse_location(message)
    if not loc:
        await message.answer(tr(('new_poi', 'no_location')),
                             reply_markup=location_keyboard())
        return
    if not valid_location(loc):
        await message.answer(tr(('new_poi', 'location_out')),
                             reply_markup=location_keyboard())
        return
    await state.update_data(lon=loc.lon, lat=loc.lat)
    await EditState.keywords.set()
    await message.answer(tr(('new_poi', 'keywords')), reply_markup=cancel_keyboard())

@dp.message_handler(state=EditState.keywords)
async def new_keywords(message: types.Message, state: FSMContext):
    keywords = split_tokens(message.text)
    if not keywords:
        await message.answer(tr(('new_poi', 'no_keywords')))
        return
    # Create a POI
    data = await state.get_data()
    poi = POI(
        name=data['name'],
        location=Location(lat=data['lat'], lon=data['lon']),
        keywords=' '.join(keywords)
    )
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(message.from_user, state, tr(('new_poi', 'confirm')))

def format(v, yes=None, no=None, null=None):
    if v is None or v == '':
        return f'{null or tr(("editor", "unknown"))}'
    if isinstance(v, str):
        return h(v)
    if isinstance(v, bool):
        return (yes or tr(('editor', 'bool_yes'))) if v else (no or tr(('editor', 'bool_no')))
    if isinstance(v, (int, float)):
        return v
    if isinstance(v, hoh.OHParser):
        return h(v.field)
    if isinstance(v, Location):
        return f'v.lat, v.lon'
    return str(v)

def new_keyboard():
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton('💾 ' + tr('save'), callback_data='save'),
        types.InlineKeyboardButton('❌ ' + tr('cancel'), callback_data='cancel')
    )

async def print_edit_options(user: types.User, state: FSMContext, comment=None):
    poi = (await state.get_data())['poi']
    lines = []
    m = tr(('editor', 'panel'))
    lines.append(f'{format(poi.name)}')
    lines.append('')
    lines.append(f'/edesc {m["desc"]}: {format(poi.description, null=m["none"])}')
    lines.append(f'/ekey {m["keywords"]}: {format(poi.keywords)}')
    lines.append(f'/etag {m["tag"]}: {format(poi.tag)}')
    lines.append(f'/ehouse {m["house"]}: {format(poi.house_name)}')
    lines.append(f'/efloor {m["floor"]}: {format(poi.floor)}')
    lines.append(f'/eaddr {m["addr"]}: {format(poi.address_part)}')
    lines.append(f'/ehour {m["hours"]}: {format(poi.hours_src)}')
    lines.append(f'/eloc {m["loc"]}: '
                 ''
                 f'{m["loc_browse"]}')
    lines.append(f'/ephone {m["phone"]}: {format("; ".join(poi.phones))}')
    lines.append(f'/ewifi {m["wifi"]}: {format(poi.has_wifi)}')
    lines.append(f'/ecard {m["card"]}: {format(poi.accepts_cards)}')
    if poi.links:
        links = ', '.join([f'{h(l[0])}' for l in poi.links])
    else:
        links = f'{m["none"]}'
    lines.append(f'/elink {m["links"]}: {links}')
    lines.append(f'/ecom {m["comment"]}: {format(poi.comment, null=m["none"])}')
    if poi.photo_out and poi.photo_in:
        photos = m['photo_both']
    elif poi.photo_out:
        photos = m['photo_out']
    elif poi.photo_in:
        photos = m['photo_in']
    else:
        photos = m['none']
    lines.append(f'{m["photo"]}: {photos} ({m["photo_comment"]})')
    if poi.id:
        if poi.delete_reason:
            lines.append(f'{m["deleted"]}: {format(poi.delete_reason)}. '
                         f'{m["restore"]}: /undelete')
        else:
            lines.append(f'🗑 {m["delete"]}: /delete')
        lines.append(f'✉ {m["msg"]}: /msg')

    content = '\n'.join(lines)
    if comment is None:
        comment = tr(('new_poi', 'confirm2'))
    if comment:
        content += '\n\n' + h(comment)
    reply = await bot.send_message(
        user.id, content, parse_mode=HTML, reply_markup=new_keyboard(),
        disable_web_page_preview=True)
    await state.update_data(reply=reply.message_id)

def cancel_attr_kbd():
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton(tr(('editor', 'cancel')), callback_data='cancel_attr')
    )

def edit_loc_kbd(poi):
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton(
            tr(('editor', 'latlon')),
            url='https://zverik.github.io/latlon/#18/'
                f'{poi.location.lat}/{poi.location.lon}"'),
        types.InlineKeyboardButton(tr(('editor', 'cancel')), callback_data='cancel_attr')
    )

def boolean_kbd(attr: str):
    return types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton(tr(('editor', 'bool_true')),
                                   callback_data=BOOL_CB.new(attr=attr, value='true')),
        types.InlineKeyboardButton(tr(('editor', 'bool_false')),
                                   callback_data=BOOL_CB.new(attr=attr, value='false')),
        types.InlineKeyboardButton(tr(('editor', 'bool_none')),
                                   callback_data=BOOL_CB.new(attr=attr, value='null')),
        types.InlineKeyboardButton(tr(('editor', 'cancel')), callback_data='cancel_attr')
    )

def tag_kbd(page: int = 1):
    ROWS = 5
    tags = config.TAGS['suggest_tags']
    kbd = types.InlineKeyboardMarkup(row_width=3)
    if (page - 1) * ROWS * 3 >= len(tags):
        page = 1
    for tag in tags[(page - 1) * ROWS * 3:page * ROWS * 3]:
        kbd.insert(types.InlineKeyboardButton(config.TAGS['tags'].get(tag, [tag])[0],
                                              callback_data=TAG_CB.new(tag=tag)))
    kbd.add(
        types.InlineKeyboardButton(tr(('editor', 'cancel')), callback_data='cancel_attr'),
        types.InlineKeyboardButton(tr(('editor', 'next_page')) + ' ⏭',
                                   callback_data=TAG_PAGE_CB.new(page=str(page + 1)))
    )
    return kbd

@dp.callback_query_handler(TAG_PAGE_CB.filter(), state=EditState.attr)
async def next_page(query: types.CallbackQuery, callback_data: Dict[str, str]):
    kbd = tag_kbd(int(callback_data['page']))
    await bot.edit_message_reply_markup(
        query.from_user.id, query.message.message_id, reply_markup=kbd)

async def delete_msg(source: Union[types.Message, types.CallbackQuery],
                     message_id: Union[int, FSMContext] = None):
    user_id = source.from_user.id
    if isinstance(message_id, FSMContext):
        message_id = (await message_id.get_data()).get('reply')
    if isinstance(source, types.CallbackQuery):
        if isinstance(message_id, list):
            message_id.append(source.message.message_id)
        else:
            message_id = source.message.message_id

    if message_id:
        if not isinstance(message_id, list):
            message_id = [message_id]
        for msg_id in message_id:
            if msg_id:
                try:
                    await bot.delete_message(user_id, msg_id)
                except MessageToDeleteNotFound:
                    pass
                except TelegramAPIError:
                    logging.exception('Failed to delete bot message')

@dp.message_handler(commands='ephoto', state=EditState.confirm)
async def show_photos(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    for photo, where in [(poi.photo_out, 'photo_out'), (poi.photo_in, 'photo_in')]:
        if photo:
            path = os.path.join(config.PHOTOS, photo + '.jpg')
            if os.path.exists(path):
                kbd = types.InlineKeyboardMarkup().add(
                    types.InlineKeyboardButton(
                        '🗑 ' + tr(('editor', 'photo_del')),
                        callback_data=PHOTO_CB.new(name=photo, which='unlink')),
                    types.InlineKeyboardButton(
                        tr(('editor', 'cancel')), callback_data='cancel_attr')
                )
                await message.answer_photo(
                    types.InputFile(path), caption=tr(('editor', where)), reply_markup=kbd)

@dp.message_handler(commands='eout', state=EditState.confirm)
async def suggest_photo_out(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    pois = await db.get_poi_around(poi.location, 10, floor=poi.floor)
    photos = [p.photo_out for p in pois if p.photo_out]
    if not photos:
        await message.answer(tr(('editor', 'no_photos_around')))
        return

    photo_dist = {
        name: min(poi.location.distance(p.location) for p in pois if p.photo_out == name)
        for name in photos
    }
    photo_cnt = Counter(photos)
    photos = sorted(photo_cnt, key=lambda p: (int(photo_dist.get(p, 1000) / 10),
                                              100 - photo_cnt[p]))
    photos = photos[:3]

    kbd = types.InlineKeyboardMarkup(row_width=5)
    media = types.MediaGroup()
    for i, photo in enumerate(photos, 1):
        path = os.path.join(config.PHOTOS, photo + '.jpg')
        if os.path.exists(path):
            file_ids = await db.find_file_ids({photo: os.path.getsize(path)})
            if photo in file_ids:
                media.attach_photo(file_ids[photo])
            else:
                media.attach_photo(types.InputFile(path))
            kbd.insert(types.InlineKeyboardButton(
                str(i), callback_data=PHOTO_CB.new(name=photo, which='out')))
    kbd.insert(types.InlineKeyboardButton(
        tr(('editor', 'cancel')), callback_data='cancel_attr'))
    await delete_msg(message, state)
    msg = await bot.send_media_group(message.from_user.id, media=media)
    msg2 = await message.answer(tr(('editor', 'choose_photo')), reply_markup=kbd)
    replies = msg.message_id if not isinstance(msg, list) else [m.message_id for m in msg]
    replies.append(msg2.message_id)
    await state.update_data(reply=replies)

@dp.message_handler(state=EditState.confirm, content_types=types.ContentType.PHOTO)
async def upload_photo(message: types.Message, state: FSMContext):
    file_id = message.photo[-1].file_id
    name = await db.find_path_for_file_id(file_id)
    downloaded = False
    if not name:
        try:
            f = await bot.get_file(file_id)
            name = (''.join(random.sample(ascii_lowercase, 4)) +
                    datetime.now().strftime('%y%m%d%H%M%S'))
            path = os.path.join(config.PHOTOS, name + '.jpg')
            await f.download(path)
        except TelegramAPIError:
            logging.exception('Image upload fail')
            await message.answer(tr(('editor', 'upload_fail')))
            return
        if not os.path.exists(path):
            await message.answer(tr(('editor', 'upload_fail')))
            return
        await db.store_file_id(name, os.path.getsize(path), file_id)
        downloaded = True

    kbd = types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton(tr(('editor', 'photo_out')),
                                   callback_data=PHOTO_CB.new(name=name, which='out')),
        types.InlineKeyboardButton(tr(('editor', 'photo_in')),
                                   callback_data=PHOTO_CB.new(name=name, which='in')),
        types.InlineKeyboardButton(
            '🗑 ' + tr(('editor', 'photo_del')),
            callback_data=PHOTO_CB.new(name=name, which='del' if downloaded else 'skip'))
    )
    await delete_msg(message, state)
    await message.answer(tr(('editor', 'photo')), reply_markup=kbd)

@dp.callback_query_handler(PHOTO_CB.filter(), state=EditState.confirm)
async def store_photo(query: types.CallbackQuery, callback_data: Dict[str, str],
                      state: FSMContext):
    poi = (await state.get_data())['poi']
    name = callback_data['name']
    path = os.path.join(config.PHOTOS, name + '.jpg')
    if not os.path.exists(path):
        await query.answer(tr(('editor', 'photo_lost')))
        return
    which = callback_data['which']
    if which == 'out':
        poi.photo_out = name
    elif which == 'in':
        poi.photo_in = name
    elif which == 'unlink':
        if poi.photo_out == name:
            poi.photo_out = None
        elif poi.photo_in == name:
            poi.photo_in = None
    elif which == 'del':
        os.remove(path)
        await query.answer(tr(('editor', 'photo_deleted')))
    else:
        await query.answer(tr(('editor', 'photo_forgot')))
    await delete_msg(query)
    await state.set_data({'poi': poi})
    await print_edit_options(query.from_user, state)

@dp.message_handler(commands='msg', state=EditState.confirm)
async def message_intro(message: types.Message, state: FSMContext):
    user = await get_user(message.from_user)
    if user.is_moderator():
        await message.answer(tr(('editor', 'cant_message')))
        return
    await delete_msg(message, state)
    reply = await message.answer(
        tr(('editor', 'message')), reply_markup=cancel_attr_kbd())
    await EditState.message.set()
    await state.update_data(reply=reply.message_id)

async def print_edit_message(message: types.Message, state: FSMContext,
                             attr: str, dash: bool = False, poi_attr: str = None,
                             msg_attr: str = None, kbd=None, value='-',
                             content: str = None):
    reply0 = None
    if value is not None:
        poi = (await state.get_data())['poi']
        if value == '-':
            pvalue = getattr(poi, poi_attr or attr)
        else:
            pvalue = value(poi)
        if pvalue:
            reply0 = (await message.answer(str(pvalue))).message_id

    if not content:
        content = tr(('editor', msg_attr or attr))
        if dash:
            content += ' ' + tr(('editor', 'dash'))
    if not kbd:
        kbd = cancel_attr_kbd()
    await delete_msg(message, state)
    reply = await message.answer(content, reply_markup=kbd, disable_web_page_preview=True)
    await EditState.attr.set()
    await state.update_data(attr=attr, reply=[reply.message_id, reply0])

@dp.message_handler(commands='ename', state=EditState.confirm)
async def edit_name(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'name')

@dp.message_handler(commands='edesc', state=EditState.confirm)
async def edit_desc(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'desc', dash=True, poi_attr='description')

@dp.message_handler(commands='etag', state=EditState.confirm)
async def edit_tag(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'tag', dash=True, kbd=tag_kbd())

@dp.message_handler(commands='ecom', state=EditState.confirm)
async def edit_comment(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'comment', dash=True)

@dp.message_handler(commands='ekey', state=EditState.confirm)
async def edit_keywords(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'keywords')

@dp.message_handler(commands='eaddr', state=EditState.confirm)
async def edit_address(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'address', dash=True, poi_attr='address_part')

@dp.message_handler(commands='eloc', state=EditState.confirm)
async def edit_location(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    await print_edit_message(message, state, 'location', value=None,
                             kbd=edit_loc_kbd(poi))

@dp.message_handler(commands='ephone', state=EditState.confirm)
async def edit_phones(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'phones', dash=True,
                             value=lambda p: '; '.join(p.phones))

@dp.message_handler(commands='ehour', state=EditState.confirm)
async def edit_hours(message: types.Message, state: FSMContext):
    await print_edit_message(message, state, 'hours', dash=True,
                             poi_attr='hours_src')

@dp.message_handler(commands='ehouse', state=EditState.confirm)
async def edit_house(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    houses = await db.get_houses()
    houses.sort(key=lambda h: poi.location.distance(h.location))
    houses = houses[:3]

    # Prepare the map
    map_file = get_map([h.location for h in houses], ref=poi.location)
    # Prepare the keyboard
    kbd = types.InlineKeyboardMarkup(row_width=1)
    for i, house in enumerate(houses, 1):
        prefix = '✅ ' if house == poi.house else ''
        kbd.add(types.InlineKeyboardButton(
            f'{prefix} {i} {house.name}', callback_data=HOUSE_CB.new(hid=house.key)))
    kbd.add(types.InlineKeyboardButton(
        tr(('editor', 'cancel')), callback_data='cancel_attr'))

    # Finally send the reply
    await delete_msg(message, state)
    if map_file:
        await message.answer_photo(types.InputFile(map_file.name),
                                   caption=tr(('editor', 'house')), reply_markup=kbd)
        map_file.close()
    else:
        await message.answer(tr(('editor', 'house')), reply_markup=kbd)

@dp.callback_query_handler(HOUSE_CB.filter(), state=EditState.confirm)
async def update_house(query: types.CallbackQuery, callback_data: Dict[str, str],
                       state: FSMContext):
    poi = (await state.get_data())['poi']
    hid = callback_data['hid']
    poi.house = hid
    h_data = await db.get_poi_by_key(hid)
    if h_data:
        poi.house_name = h_data.name
    await delete_msg(query)
    await state.set_data({'poi': poi})
    await print_edit_options(query.from_user, state)

@dp.message_handler(commands='efloor', state=EditState.confirm)
async def edit_floor(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    floors = await db.get_floors_by_house(poi.house)
    if floors and floors != [None]:
        kbd = types.InlineKeyboardMarkup(row_width=3)
        for floor in floors:
            if floor is not None:
                kbd.insert(types.InlineKeyboardButton(
                    floor, callback_data=FLOOR_CB.new(floor=floor)))
        kbd.insert(types.InlineKeyboardButton(
            tr(('editor', 'cancel')), callback_data='cancel_attr'))
    else:
        kbd = cancel_attr_kbd()
    await print_edit_message(message, state, 'floor', dash=True,
                             kbd=kbd, value=None)

@dp.callback_query_handler(FLOOR_CB.filter(), state=EditState.attr)
async def update_floor(query: types.CallbackQuery, callback_data: Dict[str, str],
                       state: FSMContext):
    poi = (await state.get_data())['poi']
    floor = callback_data['floor']
    poi.floor = floor if floor != '-' else None
    await delete_msg(query)
    await EditState.confirm.set()
    await state.set_data({'poi': poi})
    await print_edit_options(query.from_user, state)

@dp.message_handler(commands='ewifi', state=EditState.confirm)
async def edit_wifi(message: types.Message, state: FSMContext):
    await delete_msg(message, state)
    reply = await message.answer(
        tr(('editor', 'wifi')), reply_markup=boolean_kbd('wifi'))
    await state.update_data(reply=reply.message_id)

@dp.message_handler(commands='ecard', state=EditState.confirm)
async def edit_cards(message: types.Message, state: FSMContext):
    await delete_msg(message, state)
    reply = await message.answer(
        tr(('editor', 'cards')), reply_markup=boolean_kbd('cards'))
    await state.update_data(reply=reply.message_id)

@dp.callback_query_handler(BOOL_CB.filter(), state=EditState.confirm)
async def update_boolean(query: types.CallbackQuery, callback_data: Dict[str, str],
                         state: FSMContext):
    poi = (await state.get_data())['poi']
    attr = callback_data['attr']
    svalue = callback_data['value']
    if svalue == 'null':
        value = None
    else:
        value = svalue == 'true'
    if attr == 'wifi':
        poi.has_wifi = value
    elif attr == 'cards':
        poi.accepts_cards = value
    else:
        query.answer(f'Unknown field {attr}')
    await delete_msg(query)
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(query.from_user, state)

@dp.callback_query_handler(TAG_CB.filter(), state=EditState.attr)
async def update_tag(query: types.CallbackQuery, callback_data: Dict[str, str],
                     state: FSMContext):
    poi = (await state.get_data())['poi']
    poi.tag = callback_data['tag']
    await delete_msg(query)
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(query.from_user, state)

@dp.message_handler(commands='elink', state=EditState.confirm)
async def edit_links(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    if poi.links:
        content = tr(('editor', 'links_have')) + '\n\n'
        content += '\n'.join([f'🔗 {h(l[0])}: {h(l[1])}' for l in poi.links])
    else:
        content = tr(('editor', 'no_links'))
    content += '\n\n' + tr(('editor', 'links'))
    await print_edit_message(message, state, 'links', value=None, content=content)

@dp.message_handler(commands='delete', state=EditState.confirm)
async def delete_poi_prompt(message: types.Message, state: FSMContext):
    poi = (await state.get_data())['poi']
    if poi.delete_reason:
        user = await get_user(message.from_user)
        if not user.is_moderator():
            await message.answer(tr(('editor', 'delete_twice')))
        else:
            await db.delete_poi_forever(user.id, poi)
            await state.finish()
            await message.answer(tr(('editor', 'deleted2')), reply_markup=get_buttons())
    else:
        await message.answer(tr(('editor', 'delete')), reply_markup=cancel_attr_kbd())
        await EditState.attr.set()
        await state.update_data(attr='delete')

@dp.message_handler(commands='undelete', state=EditState.confirm)
async def undelete_poi(message: types.Message, state: FSMContext):
    user = await get_user(message.from_user)
    if not user.is_moderator():
        await message.answer(tr(('editor', 'cant_restore')))
        return
    poi = (await state.get_data())['poi']
    await db.restore_poi(message.from_user.id, poi)
    await state.finish()
    if user.review:
        kbd = types.InlineKeyboardMarkup().add(types.InlineKeyboardButton(
            '🗒 ' + tr(('review', 'continue')), callback_data='continue_review'))
    else:
        kbd = get_buttons()
    await message.answer(tr(('editor', 'restored')), reply_markup=kbd)

@dp.callback_query_handler(text='cancel_attr', state=EditState.all_states)
async def cancel_attr(query: types.CallbackQuery, state: FSMContext):
    await delete_msg(query, state)
    await EditState.confirm.set()
    await print_edit_options(query.from_user, state)

@dp.message_handler(content_types=types.ContentType.LOCATION, state=EditState.attr)
async def store_location(message: types.Message, state: FSMContext):
    data = await state.get_data()
    poi = data['poi']
    attr = data['attr']
    if attr != 'location':
        return
    loc = parse_location(message)
    if not loc:
        await message.answer(tr(('new_poi', 'no_location')), reply_markup=edit_loc_kbd(poi))
        return
    poi.location = loc
    await delete_msg(message, state)
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(message.from_user, state)

RE_URL = re.compile(r'^https?://')
# TODO: use hours_abbr and DOW
RE_HOURS = re.compile(r'^(?:(пн|вт|ср|чт|пт|сб|вс|mo|tu|we|th|fr|sa|su)(?:\s*-?'
                      r'\s*(пн|вт|ср|чт|пт|сб|вс|mo|tu|we|th|fr|sa|su))?\s+)?'
                      r'(\d\d?(?:[:.]\d\d)?)\s*-\s*(\d\d(?:[:.]\d\d)?)'
                      r'(?:\s+об?е?д?\s+(\d\d?(?:[:.]\d\d)?)\s*-\s*(\d\d(?:[:.]\d\d)?))?$')

def parse_hours(s):
    def norm_hour(h):
        if not h:
            return None
        if len(h) < 4:
            h += ':00'
        return h.replace('.', ':').rjust(5, '0')

    if s in ('24', '24/7'):
        return '24/7'

    HOURS_WEEK = {tr(('editor', 'hours_abbr'))[i]: DOW[i] for i in range(7)}
    HOURS_WEEK.update({DOW[i].lower(): DOW[i] for i in range(7)})

    parts = []
    for part in s.split(','):
        m = RE_HOURS.match(part.strip().lower())
        if not m:
            raise ValueError(part)
        wd = 'Mo-Su' if not m.group(1) else HOURS_WEEK[m.group(1)]
        if m.group(2):
            wd += '-' + HOURS_WEEK[m.group(2)]
        h1 = norm_hour(m.group(3))
        h2 = norm_hour(m.group(4))
        l1 = norm_hour(m.group(5))
        l2 = norm_hour(m.group(6))
        if l1:
            wd += f' {h1}-{l1},{l2}-{h2}'
        else:
            wd += f' {h1}-{h2}'
        parts.append(wd)
    return '; '.join(parts)

def parse_link(value):
    parts = value.lower().replace('. ', '.').split(None, 1)
    if len(parts) == 1 and '.' not in parts[0]:
        return parts
    if len(parts) == 1:
        parts = [tr('default_link'), parts[0]]
    REPLACE_TITLE = tr(('editor', 'link_replace'))
    REPLACE_TITLE['instagram'] = tr(('editor', 'instagram'))
    if parts[0] in REPLACE_TITLE:
        parts[0] = REPLACE_TITLE[parts[0]]
    if parts[0] == tr(('editor', 'instagram')) and 'instagram.' not in parts[1]:
        parts[1] = 'https://instagram.com/' + parts[1]
    elif parts[0] == 'vk' and 'vk.' not in parts[1]:
        parts[1] = 'https://vk.com/' + parts[1]
    if '://' not in parts[1]:
        parts[1] = 'https://' + parts[1]
    return parts

@dp.message_handler(state=EditState.attr)
async def store_attr(message: types.Message, state: FSMContext):
    data = await state.get_data()
    poi = data['poi']
    attr = data['attr']
    value = message.text.strip()

    if attr == 'name':
        if value == '-':
            await message.answer(tr(('editor', 'empty_name')), reply_markup=cancel_attr_kbd())
            return
        poi.name = value
    elif attr == 'desc':
        poi.description = None if value == '-' else value
    elif attr == 'comment':
        poi.comment = None if value == '-' else value
    elif attr == 'floor':
        poi.floor = None if value == '-' else value
    elif attr == 'tag':
        if value == '-':
            poi.tag = None
        else:
            parts = [p.strip() for p in re.split(r'[ =]+', value.lower().replace('-', '_'))]
            if len(parts) != 2 or not re.match(r'^[a-z]+$', parts[0]):
                await message.answer(tr(('editor', 'tag_format'), value),
                                     reply_markup=cancel_attr_kbd())
                return
            poi.tag = '='.join(parts)
    elif attr == 'keywords':
        new_kw = split_tokens(value)
        if new_kw:
            old_kw = [] if not poi.keywords else poi.keywords.split()
            poi.keywords = ' '.join(old_kw + new_kw)
    elif attr == 'address':
        poi.address_part = None if value == '-' else value
    elif attr == 'location':
        loc = parse_location(message)
        if not loc:
            await message.answer(tr(('new_poi', 'no_location')),
                                 reply_markup=edit_loc_kbd(poi))
            return
        if not valid_location(loc):
            await message.answer(tr(('new_poi', 'location_out')),
                                 reply_markup=edit_loc_kbd(poi))
            return
        poi.location = loc
    elif attr == 'hours':
        if value == '-':
            poi.hours = None
            poi.hours_src = None
        else:
            try:
                hours = parse_hours(value)
            except ValueError as e:
                await message.answer(tr(('editor', 'hours_format'), e),
                                     reply_markup=cancel_attr_kbd())
                return
            poi.hours_src = hours
            poi.hours = hoh.OHParser(hours)
    elif attr == 'phones':
        if not value or value == '-':
            poi.phones = []
        else:
            poi.phones = [p.strip() for p in re.split(r'[;,]', value)]
    elif attr == 'links':
        if value:
            parts = parse_link(value)
            if parts:
                if len(parts) == 1:
                    poi.links = [l for l in poi.links if l[0] != parts[0]]
                else:
                    found = False
                    for i, l in enumerate(poi.links):
                        if l[0] == parts[0]:
                            found = True
                            l[1] = parts[1]
                    if not found:
                        poi.links.append(parts)
    elif attr == 'delete':
        await db.delete_poi(message.from_user.id, poi, value)
        await state.finish()
        user = await get_user(message.from_user)
        if user.review:
            kbd = types.InlineKeyboardMarkup().add(types.InlineKeyboardButton(
                '🗒 ' + tr(('review', 'continue')), callback_data='continue_review'))
        else:
            kbd = get_buttons()
        await message.answer(tr(('editor', 'deleted')), reply_markup=kbd)
        await broadcast_str(tr(('editor', 'just_deleted'), id=poi.id, reason=value),
                            message.from_user.id)
        return
    else:
        await message.answer(tr(('editor', 'wrong_attr'), attr))

    await delete_msg(message, state)
    await state.set_data({'poi': poi})
    await EditState.confirm.set()
    await print_edit_options(message.from_user, state)

@dp.message_handler(state=EditState.confirm)
async def other_msg(message: types.Message, state: FSMContext):
    await message.answer(tr(('editor', 'other_msg')))

@dp.message_handler(state=EditState.message)
async def send_message(message: types.Message, state: FSMContext):
    user = await get_user(message.from_user)
    poi = (await state.get_data())['poi']
    if poi.id is None:
        # Won't add an abstract message to the queue
        await broadcast(message)
    else:
        await db.add_to_queue(user, poi, message.text)
    await state.finish()
    await message.answer(tr(('editor', 'msg_sent')))

@dp.callback_query_handler(state=EditState.confirm, text='save')
async def new_save(query: types.CallbackQuery, state: FSMContext):
    poi = (await state.get_data())['poi']

    # If not a moderator, mark this as needs check
    user = await get_user(query.from_user)
    if not user.is_moderator() and poi.id is None:
        poi.needs_check = True

    # Send the POI to the database
    poi_id = poi.id
    try:
        if user.is_moderator() or poi.id is None:
            poi_id = await db.insert_poi(query.from_user.id, poi)
            saved = 'saved'
            if not user.is_moderator():
                await broadcast_str(tr(('editor', 'just_added'), id=poi.id, name=poi.name),
                                    query.from_user.id)
        else:
            await db.add_to_queue(user, poi)
            await broadcast_str(tr(('queue', 'added')))
            saved = 'sent'
    except DatabaseError as e:
        await bot.send_message(
            query.from_user.id, tr(('editor', 'error_save'), e), reply_markup=new_keyboard())
        return

    # Reset state and thank the user
    await delete_msg(query, state)
    await state.finish()
    kbd = types.InlineKeyboardMarkup().add(
        types.InlineKeyboardButton('👀 ' + tr(('editor', 'saved_look')),
                                   callback_data=POI_LIST_CB.new(id=poi_id)),
        types.InlineKeyboardButton('➕ ' + tr(('editor', 'saved_add')), callback_data='new')
    )
    if user.review:
        kbd.insert(types.InlineKeyboardButton(
            '🗒 ' + tr(('review', 'continue')), callback_data='continue_review'))
    await bot.send_message(query.from_user.id, tr(('editor', saved)), reply_markup=kbd)


from __future__ import print_function

__author__ = "Nestor Bermudez"
__license__ = "MIT"
__version__ = "1.0.0"
__email__ = "nab6@illinois.edu"
__status__ = "Development"

import json
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import power_divergence, chi2_contingency
from sklearn.model_selection import train_test_split

plt.style.use('seaborn-white')
sns.set_palette('dark')

def load_brackets(fmt='TTT'):
    with open('allBrackets{}.json'.format(fmt)) as f:
        brackets = json.load(f)['brackets']
        return brackets

def triplets_f4(brackets, year):
    vectors = [list(bracket['bracket']['fullvector'])
               for bracket in brackets
               if int(bracket['bracket']['year']) < year]
    vectors = np.array(vectors, dtype=int)
    vectors = vectors[:, :60].reshape(-1, 15)
    return vectors[:, -3:]

def triplets_e8(brackets, year):
    vectors = [list(bracket['bracket']['fullvector'])
               for bracket in brackets
               if int(bracket['bracket']['year']) < year]
    vectors = np.array(vectors, dtype=int)
    vectors = vectors[:, :60].reshape(-1, 15)
    return vectors[:, [8, 9, 12]], vectors[:, [10, 11, 13]]

def triplets_r2(brackets, year):
    vectors = [list(bracket['bracket']['fullvector'])
               for bracket in brackets
               if int(bracket['bracket']['year']) < year]
    vectors = np.array(vectors, dtype=int)
    vectors = vectors[:, :60].reshape(-1, 15)
    return vectors[:, [0, 1, 8]], vectors[:, [2, 3, 9]], vectors[:, [4, 5, 10]], vectors[:, [6, 7, 11]]

def all_years_f4_triplet_plot(brackets, output_dir):
    fig, ax = plt.subplots(2, 3, figsize=(13, 7))
    for n, year in enumerate(range(2014, 2020)):
        triplets = triplets_f4(brackets, year)
        triplets = np.apply_along_axis(''.join, 1, triplets.astype(str))
        triplets, counts = np.unique(triplets, return_counts=True)

        i = n // 3
        j = n % 3
        sns.barplot(triplets, counts, color='black', ax=ax[i, j],
                    order=triplets[np.argsort(counts)])
        ax[i, j].set_title('1985 - {}'.format(year - 1))
    fig.suptitle('E8-F4 triplet dist', fontsize=16)
    # plt.show()
    plt.savefig(output_dir + '/f4_triplet_all.png')
    plt.cla()
    plt.clf()

def per_year_f4_triplet_plot(brackets, output_dir):
    for year in range(2013, 2020):
        triplets = triplets_f4(brackets, year)
        triplets = np.apply_along_axis(''.join, 1, triplets.astype(str))
        triplets, counts = np.unique(triplets, return_counts=True)
        sns.barplot(triplets, counts, color='black', order=triplets[np.argsort(counts)])
        plt.title('E8-F4 triplet distribution (1985 - {})'.format(year - 1))
        # plt.xlabel('Triplet (TTT format)')
        # plt.show()
        plt.savefig(output_dir + '/f4_triplet_{}.png'.format(year))
        plt.cla()
        plt.clf()

def check_triplet_dependency_all_data(raw_triplets):
    triplets_for_estimation = raw_triplets
    triplets_for_independence = raw_triplets
    return _check_triplet_dependency(triplets_for_estimation, triplets_for_independence)

def check_triplet_dependency_50_split(raw_triplets):
    triplets_for_estimation = raw_triplets[:len(raw_triplets) / 2, :]
    triplets_for_independence = raw_triplets[len(raw_triplets) / 2:, :]
    return _check_triplet_dependency(triplets_for_estimation, triplets_for_independence)

def check_triplet_dependency_random_split(raw_triplets):
    values = []
    for i in range(1000):
        triplets_for_estimation, triplets_for_independence = train_test_split(raw_triplets, train_size=0.5, test_size=0.5)
        p_value = _check_triplet_dependency(triplets_for_estimation, triplets_for_independence)
        values.append(p_value)
    p_value = 1000 - np.count_nonzero(np.greater(values, 0.05))
    print('# rejects', p_value)
    return p_value

def _check_triplet_dependency(triplets_for_estimation, triplets_for_independence):
    n = triplets_for_estimation.shape[0]
    triplets = np.apply_along_axis(''.join, 1,
                                   triplets_for_estimation.astype(str))
    triplets, counts = np.unique(triplets, return_counts=True)
    triplet_probs = 1. * counts / counts.sum()

    bits, counts_0 = np.unique(triplets_for_independence[:, 0], return_counts=True)
    if bits[0] != 0:
        bits = np.insert(bits, 0, 0)
        counts_0 = np.insert(counts_0, 0, 0)
    if bits[1] != 1:
        counts_0 = np.insert(counts_0, 1, 0)
    probs_0 = 1. * counts_0 / counts_0.sum()

    bits, counts_1 = np.unique(triplets_for_independence[:, 1], return_counts=True)
    if bits[0] != 0:
        bits = np.insert(bits, 0, 0)
        counts_1 = np.insert(counts_1, 0, 0)
    if bits[1] != 1:
        counts_1 = np.insert(counts_1, 1, 0)
    probs_1 = 1. * counts_1 / counts_1.sum()

    bits, counts_2 = np.unique(triplets_for_independence[:, 2], return_counts=True)
    if bits[0] != 0:
        bits = np.insert(bits, 0, 0)
        counts_2 = np.insert(counts_2, 0, 0)
    if bits[1] != 1:
        counts_2 = np.insert(counts_2, 1, 0)
    probs_2 = 1. * counts_2 / counts_2.sum()

    observed = []
    expected = []
    for t in ['000', '001', '010', '011', '100', '101', '110', '111']:
        if t not in triplets:
            triplets = np.append(triplets, t)
            triplet_probs = np.append(triplet_probs, 0.)
            counts = np.append(counts, 0.)

    # sns.barplot(triplets, counts, color='black')
    # plt.xlabel('Triplet')
    # plt.ylabel('Count')
    # plt.show()

    for i in range(8):
        triplet = np.array(list(triplets[i]), dtype=int)
        n_triplet = counts[i]
        # p_triplet = triplet_probs[i]
        p_ind = probs_0[triplet[0]] * probs_1[triplet[1]] * probs_2[triplet[2]]
        n_ind = n * p_ind
        # print('Triplet: {0}, p(i,j,k) = {1:5f}, p(i)p(j)p(k) = {2:5f}. diff = {3:5f}'.format(triplets[i], p_triplet, p_ind, p_triplet - p_ind))
        if p_ind == 0:
            continue
        observed.append(n_triplet)
        expected.append(n_ind)
        # expected.append(np.ceil(n_ind).astype(int))

    # df = pd.DataFrame.from_records(np.vstack((observed, expected)).T, index=triplets, columns=['Observed', 'Expected']).astype(int)
    # print(df)

    print('m = array(c{}, dim=c(2, 2, 2))'.format(tuple(np.array(observed)[np.argsort(triplets)].astype(int).tolist())))

    chi2, p_value = power_divergence(observed, expected, ddof=3, lambda_='pearson')
    # print('Chi-square p-value', p_value, '=>',
    #       ('Indepedent' if p_value > 0.05 else 'There is dependency'))
    # print()
    return p_value

def chi_square_test(triplets):
    for i in range(3):
        l = list(range(3))
        l.pop(i)
        data = triplets[:, l]
        _, c00 = np.unique(triplets[(data[:, 0] == 0) & (data[:, 1] == 0), i],
                           return_counts=True)
        _, c01 = np.unique(triplets[(data[:, 0] == 0) & (data[:, 1] == 1), i],
                           return_counts=True)
        _, c10 = np.unique(triplets[(data[:, 0] == 1) & (data[:, 1] == 0), i],
                           return_counts=True)
        _, c11 = np.unique(triplets[(data[:, 0] == 1) & (data[:, 1] == 1), i],
                           return_counts=True)
        table = {
            '00': {'0': c00[0], '1': c00[1]},
            '01': {'0': c01[0], '1': c01[1]},
            '10': {'0': c10[0], '1': c10[1]},
            '11': {'0': c11[0], '1': c11[1]}
        }
        table = pd.DataFrame.from_dict(table)
        chi2, p, dof, expected = chi2_contingency(table)
        print(table)
        print('p-value', p)

def check_all_triplet_dependency(brackets):
    independence_check = check_triplet_dependency_all_data
    for year in range(2019, 2020):
        print('===== YEAR = {} ====='.format(year))
        print('E8-F4 triplet')
        raw_triplets = triplets_f4(brackets, year)
        independence_check(raw_triplets)
        # chi_square_test(raw_triplets)

        print('Top S16-E8 triplet')
        top, bottom = triplets_e8(brackets, year)
        independence_check(top)
        # chi_square_test(top)

        print('Bottom S16-E8 triplet')
        independence_check(bottom)
        # chi_square_test(bottom)

        t1, t2, t3, t4 = triplets_r2(brackets, year)
        print('R1-R2 triplet 1')
        independence_check(t1)

        print('R1-R2 triplet 2')
        independence_check(t2)

        print('R1-R2 triplet 3')
        independence_check(t3)

        print('R1-R2 triplet 4')
        independence_check(t4)

if __name__ == '__main__':
    import os
    import sys
    fmt = sys.argv[1]
    output = sys.argv[2]

    if not os.path.exists(output):
        os.makedirs(output)

    brackets = load_brackets(fmt)
    check_all_triplet_dependency(brackets)
    # all_years_f4_triplet_plot(brackets, output)
    # per_year_f4_triplet_plot(brackets, output)


'''
	访问URL，重新收集subject地址，然后存入.csv
'''
from queue import Queue
from threading import Thread
import time
import random
import csv
from bs4 import BeautifulSoup
from tool import getHttpUa,getCursor,getResult,getPage,ChangeOrNot
from tool import editeHeader,editeCookies,writeMetrx,commHttpProxies,commHeaders,commCookies
#sltDLNotCom = 'select id,url,name from dlurl1 where (subject=" " or subject is null) and id>2810844 limit 50000' #13000 zuoyou
#import os;os.chdir('e:/Code/Python');import extractSubject_threading;extractSubject_threading.mainFunction()

path = 'E:/Code/Data/subUrl last 3.csv'
path2 = 'E:/Code/Data/subUrlSupply last 3.csv'
path3 = 'E:/Code/Data/subUrlNull last 3.csv'
expertList_path = 'E:/temp_csv/重置/subject url last 3.csv'
flag = 2 #选择http的来源文件
ipQueue = Queue()
uaQueue = Queue()
dlQueue = Queue()
htmlQueue = Queue()
infoQueue = Queue()
#import os;os.chdir('e:/Code/Python');import extractSubject_threading;extractSubject_threading.mainFunction()
global httpProxies
global headers
global cookies
class pageWorker(Thread):
    def __init__(self,ipQueue,uaQueue,dlQueue,htmlQueue):
        Thread.__init__(self)
        self.ipQueue = ipQueue
        self.uaQueue = uaQueue
        self.dlQueue = dlQueue
        self.htmlQueue = htmlQueue

    def run(self):
        while True:
           # 从队列中获取任务并扩展tuple
            httpProxies = commHttpProxies.copy()
            headers = commHeaders.copy()
            cookies = commCookies.copy()
            dl  = self.dlQueue.get()
            http = self.ipQueue.get()
            ua = self.uaQueue.get()
            httpProxies['https'] = http
            #修饰参数
            if ChangeOrNot() == True:#随机触发
                headers=editeHeader(ua,headers,dl['name'])#改变user agent
                cookies=editeCookies(cookies)
            time.sleep(random.randint(5, 20))#随机休眠
            #取出html
            html = str(getPage(dl['url'],httpProxies,headers,cookies))#取出url
            #放回
            self.ipQueue.put(http)
            self.uaQueue.put(ua)
            if html == ' ':#未获取成功，重新放入
                print('put back'+str(dl['id']))
                self.dlQueue.put(dl)
            #放入
            else:
                self.htmlQueue.put((html,dl))
            self.dlQueue.task_done()

class analysisWorker(Thread):
    def __init__(self,htmlQueue,infoQueue):
        Thread.__init__(self)
        self.htmlQueue = htmlQueue

    def run(self):
        while True:
           html,dl = self.htmlQueue.get()
           analysisPage(html,dl['id'],dl['name'])
           self.htmlQueue.task_done()

#获取 subject area的地址
def getSubject(soup,doc,id):
    #
    url = ' '
    stri = soup.find(text='See all subject areas')#找到目标字样
    if stri != None:
        try:
            url = 'http://dl.acm.org/' + stri.parent['href']
        except Exception:
            print('error:'+str(stri.parent))

    return url
#获取 subject area的地址
def getSubject2(soup,doc,id):
    #
    url = ' '
    stri = soup.find(text='See all author supplied keywords')#找到目标字样
    if stri != None:
        try:
            url = 'http://dl.acm.org/' + stri.parent['href']
        except Exception:
            print('error:'+str(stri.parent))

    return url
def analysisPage(doc,id,name):
    #
    soup = BeautifulSoup(''.join(doc),"lxml")
    subUrl = getSubject(soup,doc,id)
    if subUrl==' ':
        subUrl = getSubject2(soup,doc,id)
        if subUrl == ' ':
            #空占位
            try:
                #print([id,'Null'])
                writeMetrx(path3,[[id,'Null']])
                print('complete null: '+str(id))
            except Exception:
                print('write list error3'+str(id))
        else:
            try:
                #print([id,subUrl])
                writeMetrx(path2,[[id,name,subUrl]])
                print('complete: '+str(id))
            except Exception:
                print('write list error2'+str(id))
    #
    else:
        try:
            print([id,subUrl])
            writeMetrx(path,[[id,name,subUrl]])
            print('complete: '+str(id))
        except Exception:
            print('write list error1'+str(id))

def mainFunction():
    http,uag = getHttpUa()

    for ip in http:
        ipQueue.put(ip)
    for ua in uag:
        uaQueue.put(ua)
    for k in range(2):
        aWorker = analysisWorker(htmlQueue,infoQueue)
        aWorker.daemon = True
        aWorker.start()
    print('ok1')
    for i in range(8):
        pWorker = pageWorker(ipQueue,uaQueue,dlQueue,htmlQueue)
        pWorker.daemon = True
        pWorker.start()
    print('ok2')
    '''
    conn,cur = getCursor()
    dlList = getResult(sltDLNotCom,cur)#返回url实体的二维数组
    '''
    dlList=[]
    with open(expertList_path) as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            if int(row['id']) > 0:
                dlList.append(row)
            #break
        print('total:'+str(len(dlList)))

    for dl in dlList:
        dlQueue.put(dl)
    #cur.close()
    #conn.close()

if __name__ == '__main__':
	mainFunction()


from matplotlib import pyplot as plt
import seaborn as sns

def plot():
    df = sns.load_dataset("titanic")
    fig, axes = plt.subplots(nrows=1, ncols=3,figsize=(15, 5))

    df1=df.age
    df2=df1.dropna()

    sns.distplot(df2,ax=axes[0])
    sns.countplot(x="sex", hue="alive",data=df,ax=axes[1])
    sns.countplot(x="class", hue="alive", data=df,ax=axes[2])
    plt.show()
    return axes

if __name__=='__main__':
    plot()


import logging
import random
import time
from datetime import datetime, timezone

from django.utils.translation import ugettext as _

from alarm_backends.core.cache.key import FTA_SUB_CONVERGE_DIMENSION_LOCK_KEY
from alarm_backends.service.converge.dimension import (
    DimensionCalculator,
    DimensionHandler,
)
from alarm_backends.service.converge.tasks import run_converge
from alarm_backends.service.converge.utils import list_other_converged_instances
from bkmonitor.models.fta.action import (
    ActionInstance,
    ConvergeInstance,
    ConvergeRelation,
)
from constants.action import (
    ALL_CONVERGE_DIMENSION,
    ActionStatus,
    ConvergeStatus,
    ConvergeType,
)

logger = logging.getLogger("fta_action.converge")

class ConvergeManager(object):
    def __init__(
        self,
        converge_config,
        dimension,
        start_time,
        instance,
        instance_type=ConvergeType.ACTION,
        end_timestamp=None,
        alerts=None,
    ):
        self.alerts = alerts
        self.converge_config = converge_config
        self.instance_type = instance_type
        self.instance = instance
        self.dimension = dimension
        self.is_created = False
        self.start_time = start_time
        self.end_timestamp = end_timestamp
        self.match_alarm_id_list = []
        self.converge_instance = self.get_converge_instance(start_time)
        self.start_timestamp = int(self.start_time.timestamp())
        self.biz_converge_existed = False
        if self.converge_instance:
            # 如果存在收敛对象的开始时间是介于
            create_timestamp = int(self.converge_instance.create_time.timestamp())
            if self.start_timestamp < create_timestamp:
                self.start_timestamp = self.start_timestamp
        converged_condition = {
            condition_item["dimension"]: self.converge_config["converged_condition"].get(condition_item["dimension"])
            for condition_item in self.converge_config["condition"]
        }

        self.dimension_handler = DimensionHandler(
            self.dimension,
            converged_condition,
            self.start_timestamp,
            instance_id=self.instance.id,
            end_timestamp=self.end_timestamp,
            instance_type=self.instance_type,
            strategy_id=getattr(self.instance, "strategy_id", 0),
            converged_condition=self.converge_config["converged_condition"],
        )

    def do_converge(self):
        """
        收敛计算
        """
        if self.converge_instance:
            # 当前存在同维度的汇总内容，直接返回收敛
            logger.info(
                "%s|converge(%s) skipped because converge instance of dimension(%s｜%s) existed already",
                self.instance_type,
                self.instance.id,
                self.converge_instance.id,
                self.dimension,
            )
            return True

        self.match_alarm_id_list = self.get_related_ids()
        matched_count = len(self.match_alarm_id_list)

        if not self.converge_instance and matched_count >= int(self.converge_config["count"]):
            # create converge_instance record
            logger.info("__begin to create converge_instance by %s __", self.instance.id)
            try:
                self.create_converge_instance(self.start_time)
            except Exception as error:
                logger.exception("create converge_instance by instance(%s) failed：%s", self.instance.id, error)
                return False

            logger.info("__end create converge_instance(%s) by %s __ ", self.converge_instance.id, self.instance.id)

        if not self.converge_instance or matched_count == 0:
            logger.info("$%s no matched_count , return!!", self.instance.id)
            return False

        if (
            self.is_created
            and self.converge_config.get("sub_converge_config")
            and self.converge_config.get("need_biz_converge")
        ):
            # 如果当前收敛事件关联的动作有二级收敛，需要计算维度，并且推送到对应的队列中
            self.push_to_sub_converge_queue()

        return True

    def push_to_sub_converge_queue(self):
        """
        推送二级收敛至收敛队列
        """

        sub_converge_config = self.converge_config.get("sub_converge_config", {})
        sub_converge_config["is_enabled"] = True  # 需要二级收敛的，默认都设置is_enabled为True
        if self.is_biz_converge_existed(sub_converge_config["count"]):
            # 当前已经达到了业务汇总的条件，当前的告警不发出
            self.biz_converge_existed = True

        sub_converge_info = DimensionCalculator(
            self.converge_instance,
            ConvergeType.CONVERGE,
            converge_config=sub_converge_config,
        ).calc_sub_converge_dimension()
        task_id = run_converge.delay(
            sub_converge_config,
            self.converge_instance.id,
            ConvergeType.CONVERGE,
            sub_converge_info["converge_context"],
            alerts=self.alerts,
        )
        logger.info("push converge(%s) to converge queue, task id %s", self.converge_instance.id, task_id)

    def is_biz_converge_existed(self, matched_count):
        client = FTA_SUB_CONVERGE_DIMENSION_LOCK_KEY.client
        biz_converge_lock_key = FTA_SUB_CONVERGE_DIMENSION_LOCK_KEY.get_key(
            **self.dimension_handler.get_sub_converge_label_info()
        )
        if client.incr(biz_converge_lock_key) > matched_count:
            # 如果当前的计数器大于并发数，直接返回异常
            logger.info(
                "action(%s|%s) will be skipped because count of biz_converge_lock_key(%s) is bigger than %s, ",
                self.instance.id,
                self.converge_instance.id,
                biz_converge_lock_key,
                matched_count,
            )
            return True
        client.expire(biz_converge_lock_key, FTA_SUB_CONVERGE_DIMENSION_LOCK_KEY.ttl)
        return False

    def get_related_ids(self):
        """
        获取到当前收敛对象相关的处理动作
        """

        matched_related_ids = self.dimension_handler.get_by_condition()

        # 仅获取同一种收敛事件的id
        matched_related_ids = [
            int(related_id.split("_")[-1]) for related_id in matched_related_ids if self.instance_type in related_id
        ]

        if not matched_related_ids:
            return []

        converge_relations = ConvergeRelation.objects.filter(
            related_id__in=matched_related_ids, related_type=self.instance_type
        )
        if self.converge_instance:
            # 最后计算出来的匹配的，必须是当前收敛的对象中的关联ID或者未关联的
            converge_relations = converge_relations.exclude(converge_id=self.converge_instance.id)

        converged_related_ids = converge_relations.values_list("related_id", flat=True)

        matched_related_ids = list(set(matched_related_ids) - set(converged_related_ids))
        logger.info(
            "$%s:%s dimension alarm list: %s, %s",
            self.instance.id,
            self.instance_type,
            self.dimension,
            len(matched_related_ids),
        )
        return matched_related_ids

    def create_converge_instance(self, start_time=None):

        self.insert_converge_instance()
        if start_time and self.converge_instance.create_time < start_time:
            logger.info(
                "incident end by start_time %s (%s < %s)",
                self.converge_instance.id,
                self.converge_instance.create_time,
                start_time,
            )
            self.end_converge_by_id(self.converge_instance.id)

        return self.converge_instance

    @classmethod
    def get_fixed_dimension(cls, dimension):
        return "{} fixed at {} {}".format(dimension, int(datetime.now().timestamp()), random.randint(100, 999))

    @classmethod
    def end_converge_by_id(cls, converge_id, conv_instance=None):
        logger.info("conv_instance end by id %s", converge_id)
        if conv_instance is None:
            conv_instance = ConvergeInstance.objects.get(id=converge_id)
        if conv_instance and not conv_instance.end_time:
            conv_instance.end_time = datetime.now(tz=timezone.utc)
            conv_instance.dimension = cls.get_fixed_dimension(conv_instance.dimension)
            conv_instance.save(update_fields=["end_time", "dimension"])
        if conv_instance.converge_type == ConvergeType.CONVERGE:
            # 如果是多级收敛，需要关闭掉其他的收敛
            for conv_id in ConvergeRelation.objects.filter(converge_id=converge_id).values_list(
                "related_id", flat=True
            ):
                cls.end_converge_by_id(conv_id)
        logger.info("conv_instance %s already end at %s", converge_id, conv_instance.end_time)

    def insert_converge_instance(self):
        try:
            converged_condition_display = []
            for converged_condition_key in self.converge_config["converged_condition"]:
                if converged_condition_key == "action_id":
                    continue
                converged_condition_display.append(
                    str(ALL_CONVERGE_DIMENSION.get(converged_condition_key, converged_condition_key))
                )

            description = _("在{}分钟内，当具有相同{}的告警超过{}条以上，在执行相同的处理套餐时，进行告警防御").format(
                self.converge_config["timedelta"] // 60,
                ",".join(converged_condition_display),
                self.converge_config["count"],
            )

            self.converge_instance = ConvergeInstance.objects.create(
                converge_config=self.converge_config,
                bk_biz_id=self.instance.bk_biz_id,
                dimension=self.dimension,
                description=description,
                content=self.converge_config.get("context", "{}"),
                end_time=None,
                converge_func=self.converge_config["converge_func"],
                converge_type=self.instance_type,
                is_visible=True,
            )
        except BaseException as error:
            logger.error("insert_converge_instance error %s", str(error))
            self.is_created = False
            self.converge_instance = self.get_converge_instance()

        else:
            self.is_created = True

    def get_converge_instance(self, start_time=None):
        try:
            converge_instance = ConvergeInstance.objects.filter(dimension=self.dimension).first()
        except Exception:
            converge_instance = None
        if converge_instance and start_time and converge_instance.create_time < start_time:
            # 如果存在收敛并且不在当前收敛期的，直接关闭
            logger.info(
                "incident end by start_time %s (%s < %s)",
                converge_instance.id,
                converge_instance.create_time,
                start_time,
            )
            self.end_converge_by_id(converge_instance.id)
            converge_instance = None
        self.converge_instance = converge_instance
        return self.converge_instance

    def connect_converge(self, status=ConvergeStatus.SKIPPED):
        """关联告警"""
        try:
            is_primary = True if self.is_created else False
            if is_primary:
                converge_status = ConvergeStatus.EXECUTED
            else:
                converge_status = ConvergeStatus.SKIPPED if status else ConvergeStatus.EXECUTED

            ConvergeRelation.objects.create(
                related_id=self.instance.id,
                converge_id=self.converge_instance.id,
                related_type=self.instance_type,
                is_primary=is_primary,
                converge_status=converge_status,
                alerts=getattr(self.instance, "alerts", []),
            )
        except BaseException as error:
            # 创建失败的原因，是由于已经关联过
            logger.info("create converge relation record failed %s, is_created: %s", str(error), self.is_created)
            return

        if self.is_created and self.match_alarm_id_list:
            """ "统计恢复"""
            other_converged_instances = list_other_converged_instances(
                self.match_alarm_id_list, self.instance, self.instance_type
            )
            if self.instance_type == ConvergeType.ACTION:
                other_converged_instances = (
                    ActionInstance.objects.filter(id__in=self.match_alarm_id_list)
                    .exclude(status__in=[ActionStatus.RECEIVED, ActionStatus.SLEEP, ActionStatus.WAITING])
                    .exclude(id=self.instance.id)
                )

            if other_converged_instances.exists():
                ConvergeRelationManager.connect(
                    self.converge_instance.id,
                    set(other_converged_instances.values_list("id", flat=True)),
                    self.instance_type,
                    self.instance.id,
                    converge_status=ConvergeStatus.SKIPPED
                    if self.instance_type == ConvergeType.CONVERGE
                    else ConvergeStatus.EXECUTED,
                )
                if self.instance_type == ConvergeType.CONVERGE:
                    # 当时二级收敛的时候，需要抑制一级收敛
                    other_converged_instances.update(is_visible=False)

        if self.instance_type == ConvergeType.CONVERGE:
            self.instance.is_visible = False
            self.instance.save(update_fields=["is_visible"])

        description = self.converge_config.get("description")
        if not description or description == self.converge_instance.description:
            return

        ConvergeInstance.objects.filter(id=self.converge_instance.id).update(description=description)

    def count_instance(self):
        return ConvergeRelationManager.count(converge_id=self.converge_instance.id)

    def converge(self, conv_instance_id=None):
        # 设置收敛有效性
        hide_converge_instances, show_converge_instances = self.get_related_converge_instances(conv_instance_id)
        self.trigger_converge(hide_converge_instances, is_visible=False)
        self.trigger_converge(show_converge_instances, is_visible=True)

    def get_related_converge_instances(self, converge_id):
        converge_id = converge_id
        same_converge_instances = [converge_id]
        hide_converge_instances = []
        show_converge_instances = []
        related_ids = ConvergeRelation.objects.filter(converge_id=converge_id).values_list("related_id", flat=True)
        logger.info("$%s related_ids: %s", converge_id, related_ids)
        related_instances = ConvergeRelation.objects.filter(related_id__in=related_ids, related_type=self.instance_type)
        converge_instances = {}
        for inst in related_instances:
            converge_instances.setdefault(inst.converge_id, set()).add(inst.related_id)
        logger.info("$found converge_instances from %s: %s", converge_id, converge_instances.keys())
        if converge_id in converge_instances:
            std_alarms = converge_instances.pop(converge_id)
        else:
            std_alarms = set(related_ids)
        for other_inc, target_alarms in converge_instances.items():
            # 真子集的情况下, 不判断状态和先后, 直接隐藏子集
            if std_alarms > target_alarms:
                hide_converge_instances.append(other_inc)
                show_converge_instances.append(converge_id)
            elif std_alarms == target_alarms:
                # 全等的情况下, 放在一起全部比较，优先隐藏后面建立的
                same_converge_instances.append(other_inc)
        same_converge_instances = sorted(set(same_converge_instances))
        # 没有显示的收敛，则取全等的收敛最先建立的显示
        if same_converge_instances and not show_converge_instances:
            show_converge_instances.append(same_converge_instances[0])
            hide_converge_instances.extend(same_converge_instances[1:])
        # 否则隐藏相等的收敛
        else:
            hide_converge_instances.extend(same_converge_instances)
        logger.info("incident hide %s | show %s", set(hide_converge_instances), set(show_converge_instances))
        return hide_converge_instances, show_converge_instances

    @classmethod
    def trigger_converge(cls, converge_instances, is_visible):
        converge_instances = sorted(list(set(map(int, converge_instances))))
        if not converge_instances:
            return
        ConvergeInstance.objects.filter(id__in=converge_instances).update(is_visible=is_visible)

class ConvergeRelationManager(object):
    @staticmethod
    def count(converge_id):
        return ConvergeRelation.objects.filter_by(converge_id=converge_id).count()

    @staticmethod
    def index(converge_id, relate_id):
        related_ids = list(ConvergeRelation.objects.filter(converge_id=converge_id).values_list("relate_id", flat=True))
        return related_ids.index(relate_id)

    @staticmethod
    def connect(converge_id, related_ids, related_type, instance_id, converge_status=ConvergeStatus.SKIPPED):
        """关联收敛关系"""
        try:
            ConvergeRelationManager._connect(converge_id, related_ids, related_type, instance_id, converge_status)
        except BaseException:
            time.sleep(random.randint(1, 100) / 100.0)
            ConvergeRelationManager._connect(converge_id, related_ids, related_type, instance_id, converge_status)

    @staticmethod
    def _connect(
        converge_id,
        related_ids,
        related_type=ConvergeType.ACTION,
        instance_id=None,
        converge_status=ConvergeStatus.SKIPPED,
    ):
        """
        因为必定有很多重复记录，需要使用mysql IGNORE特性
        """
        converge_relations = [
            ConvergeRelation(
                converge_id=converge_id,
                related_id=related_id,
                related_type=related_type,
                converge_status=converge_status,
            )
            for related_id in related_ids
        ]
        ConvergeRelation.objects.ignore_blur_create(converge_relations)

        logger.info("$converge connect %s :instance_id(%s) len(%s)", converge_id, instance_id, len(related_ids))


import os

try:
    from neptune import init_run
except ImportError:
    from neptune.new import init_run

from src.neptune_detectron2 import NeptuneHook
from tests.utils import get_images

def test_e2e(cfg, trainer):

    run = init_run()

    get_images()
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    trainer.resume_or_load(resume=False)

    hook = NeptuneHook(run=run, log_checkpoints=True, log_model=True)
    trainer.register_hooks([hook])
    trainer.train()

    assert run.exists("training/config")

    assert run.exists("training/model/checkpoints/checkpoint_iter_0")

    assert run.exists("training/model/checkpoints/checkpoint_final")

    assert isinstance(run["training/model/summary"].fetch(), str)

    cls_accuracy_vals = run["training/metrics/fast_rcnn/cls_accuracy"].fetch_values()
    assert 0 < cls_accuracy_vals.iloc[-1]["value"] <= 1
    run.stop()


# coding: utf-8

# # Support Vector Regression
# ##### Using SVR to predict the MPG of vehicles

# In[2]:

import pandas as pd

# In[3]:

print(pd.__version__)

# ### Download Auto MPG data set
# Download Link: https://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data
#
# Summary: Given 8 pieces of information (features) about a vehicle, predict its mileage
#
# Notes:
# * The file does not come with headers, so we specify them explicitly

# In[4]:

auto_data = pd.read_csv('../data/auto-mpg.data', delim_whitespace = True, header=None,
                   names = ['mpg',
                            'cylinders',
                            'displacement',
                            'horsepower',
                            'weight',
                            'acceleration',
                            'model',
                            'origin',
                            'car_name'])

# In[5]:

auto_data.head()

# #### Check if the car_name feature can be helpful

# In[6]:

len(auto_data['car_name'].unique())

# In[7]:

len(auto_data['car_name'])

# #### Drop the car_name feature from the data frame
# There are too many unique values for any pattern to be detected

# In[67]:

auto_data = auto_data.drop('car_name', axis=1)
auto_data.head()

# #### Converting a numeric value for origin to something more meaningful
# * The values 1,2,3 represent America, Europe and Asia respectively
# * This is the first step before we apply one-hot-encoding for this feature
# * Renaming will give us more meaningful column names after one-hot-encoding is applied

# In[68]:

auto_data['origin'] = auto_data['origin'].replace({1: 'america', 2: 'europe', 3: 'asia'})
auto_data.head()

# #### Applying one-hot-encoding now will give us more meaningful column names

# In[69]:

auto_data = pd.get_dummies(auto_data, columns=['origin'])
auto_data.head()

# #### Convert missing values in data frame to NaN

# In[70]:

import numpy as np

auto_data = auto_data.replace('?', np.nan)

# #### Drop rows with missing values

# In[71]:

auto_data = auto_data.dropna()
auto_data

# ### Prepare training and test data
# * Define the feature vector (X) and label (Y)
# * Use train_test_split to create data subsets sets for training and validations
# * The test_size parameter specifies the proportion of the data required for testing

# In[72]:

from sklearn.model_selection import train_test_split

X = auto_data.drop('mpg', axis=1)

# Taking the labels (mpg)
Y = auto_data['mpg']

# Spliting into 80% for training set and 20% for testing set so we can see our accuracy
X_train, x_test, Y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=0)

# #### Define the Regression model
# * Use scikit-learn's SVR model
# * We start off with a linear kernel and set the regularization variable (C) to 1.0
# * The model is trained with the training data

# In[73]:

from sklearn.svm import SVR
regression_model = SVR(kernel='linear', C=1.0)
regression_model.fit(X_train, Y_train)

# #### Check the coefficients for each of the features

# In[74]:

regression_model.coef_

# #### Get R-square value with training data

# In[75]:

regression_model.score(X_train, Y_train)

# #### Use matplotlib to view the coefficients as a histogram

# In[76]:

from pandas import Series
import matplotlib.pyplot as plt
get_ipython().magic(u'matplotlib inline')

predictors = X_train.columns
coef = Series(regression_model.coef_[0],predictors).sort_values()
coef.plot(kind='bar', title='Modal Coefficients')

# #### Get predictions on test data

# In[77]:

from sklearn.metrics import mean_squared_error

y_predict = regression_model.predict(x_test)

# #### Compare the predicted and actual values of the MPG

# In[78]:

get_ipython().magic(u'pylab inline')
pylab.rcParams['figure.figsize'] = (15, 6)

plt.plot(y_predict, label='Predicted')
plt.plot(y_test.values, label='Actual')
plt.ylabel('MPG')

plt.legend()
plt.show()

# #### Get R-square value of predictions on test data

# In[79]:

regression_model.score(x_test, y_test)

# #### Calculate Mean Square Error

# In[80]:

regression_model_mse = mean_squared_error(y_predict, y_test)
regression_model_mse

# #### Root of Mean Square Error to measure degree to which our prediction is off

# In[81]:

import math

math.sqrt(regression_model_mse)

# In[ ]:

Building Machine Learning Models in Python with scikit-learn/python/m2-demo2-SVR.py


# -*- coding: utf-8 -*—
# Date: 2022/2/17 0017
# Time: 9:16
# Author:
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from sc.CIFAR10_model import HQR_CIFAR10
import time

# 加载至GPU显示
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

## 1. 加载和准备数据
train = torchvision.datasets.CIFAR10('../datasets', train=True, transform=torchvision.transforms.ToTensor(),
                                     download=True)
test = torchvision.datasets.CIFAR10('../datasets', train=False, transform=torchvision.transforms.ToTensor(),
                                    download=True)
train_len = len(train)
test_len = len(test)
print("训练数据的长度是：{}".format(train_len))
print("测试数据的长度是：{}".format(test_len))

train_data = DataLoader(train, batch_size=64, shuffle=True, num_workers=0, drop_last=True)
test_data = DataLoader(test, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

## 2. 加载网络，优化器和损失函数
learning_rate = 1e-2
HQR_model = HQR_CIFAR10()
Loss = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(HQR_model.parameters(), lr=1e-2)
# 加载至GPU进行运算
HQR_model.to(device)
Loss.to(device)

# 加载至tensorboards显示
writer = SummaryWriter('../CIFAR10_logs')

## 3. 训练和测试
# 设置全局参数

start_time = time.time()
train_step = 0
test_step = 0
for epoch in range(50):
    print("------------开始第{}轮训练-----------".format(epoch+1))
    # 开始训练
    HQR_model.train()
    for data in train_data:
        train_imgs, train_target = data
        train_imgs = train_imgs.to(device)
        train_target = train_target.to(device)
        optimizer.zero_grad()
        train_output = HQR_model(train_imgs)
        train_loss = Loss(train_output, train_target)
        train_loss.backward()
        optimizer.step()
        end_time = time.time()
        writer.add_scalar('Train', train_loss.item(), train_step)
        train_step = train_step + 1
        if train_step % 100 == 0:
            print("训练第{}次，Loss为：{}".format(train_step, train_loss.item()))
            print("训练第{}次，消耗的时间为：{}".format(train_step, (end_time - start_time)))

    # 开始测试
    HQR_model.eval()
    total_test_loss = 0
    total_test_accuracy = 0
    with torch.no_grad():
        for data in test_data:
            test_imgs, test_targets = data
            test_imgs = test_imgs.to(device)
            test_targets = test_targets.to(device)
            test_outputs = HQR_model(test_imgs)
            test_loss = Loss(test_outputs, test_targets)
            total_test_loss = total_test_loss + test_loss
            test_accuracy = (test_outputs.argmax(1) == test_targets).sum()
            total_test_accuracy = total_test_accuracy + test_accuracy

    print("测试的损失Loss为：{}".format(total_test_loss))
    writer.add_scalar('Test', total_test_loss.item(), test_step)
    print("测试的准确率为：{}".format((total_test_accuracy/test_len)))
    writer.add_scalar('ACCURACY', total_test_accuracy.item(), test_step)
    torch.save(HQR_model, '../save_model/HQR_CIFAR10_第{}轮_准确率是{}.pth'.format(test_step, (total_test_accuracy/test_len)))
    test_step = test_step + 1
    print("模型已保存！")

writer.close()


import torch
import random
from tqdm import tqdm, trange
from loss_similarity import LossSimilarity
from preprocess import init_contrastive_dataloader

@torch.no_grad()
def get_embeddings(config, memorized_samples, model, tokenizer, loss_sim: LossSimilarity,
                   before_reverse=False, return_tags=False):
    if isinstance(memorized_samples, dict):
        data_loader = memorized_samples['train_infer']
    else:
        data_loader = init_contrastive_dataloader(
            config, {'train': [], 'test': memorized_samples}, tokenizer)[0]['test']
    features, key2tags = {}, {}
    backup_alignment = model.lora_alignment
    if before_reverse:
        model.lora_alignment = None
    for batch in tqdm(data_loader, desc='embedding'):
        for key, value in batch.items():
            if isinstance(value, torch.Tensor):
                batch[key] = value.to(config['device'], non_blocking=True)
        hidden_mask = batch['input_ids'] == tokenizer.mask_token_id
        model_name = config['plm']['model_name']
        model_linear_out = model.lora_linear_out if model_name not in ['cpm', 'llama'] else model.lora_projector
        feature = loss_sim.forward_model_logits(
            '', model_linear_out, batch, model, hidden_mask, is_selector=False)[1]
        feature = feature.detach().cpu()
        assert len(feature) == len(batch['sample_keys']) == len(batch['tags'])
        for fea, sample_key, tag in zip(feature, batch['sample_keys'], batch['tags']):
            features[sample_key] = fea
            key2tags[sample_key] = int(tag)
    if before_reverse:
        model.lora_alignment = backup_alignment
    return (features, key2tags) if return_tags else features

def update_alignment_model(config, cur_embeddings: dict, mem_embeddings: dict, lora_alignment):
    # train selector from initialization
    assert config['train']['train_expert_selector'] == 0 and not config['train']['use_expert_selector']
    assert config['train']['loss_adverse_step'] <= 0

    batch_size = 32
    learning_rate = 0.0001
    epoch_num = 20
    loss_func = torch.nn.MSELoss()
    device = config['device']
    optimizer = torch.optim.Adam(lora_alignment.parameters(), lr=learning_rate)

    assert cur_embeddings.keys() == mem_embeddings.keys()
    key_set = list(cur_embeddings.keys())
    random.shuffle(key_set)
    train_size = len(key_set)

    for epoch in trange(epoch_num, desc='alignment'):
        start_idx = 0
        while start_idx < train_size:
            cur_inputs = [cur_embeddings[key] for key in key_set[start_idx:start_idx+batch_size]]
            cur_targets = [mem_embeddings[key] for key in key_set[start_idx:start_idx+batch_size]]
            cur_inputs, cur_targets = torch.stack(cur_inputs).to(device), torch.stack(cur_targets).to(device)

            optimizer.zero_grad()
            outputs = lora_alignment(cur_inputs)
            loss = loss_func(outputs, cur_targets)
            loss.backward()
            optimizer.step()

            start_idx += batch_size

    optimizer.zero_grad()


import warnings
import librosa as lr
import numpy as np
import soundfile as sf
import random as rd
import os

__all__ = ['timeStretch', 'pitchShift',
           'offsetAudio', 'fadeAudio', 'mixAudio']

descriptors = {
    "_augmentation": "base class",
    __all__[0]: "time_stretch",
    __all__[1]: "pitch_shift",
    __all__[2]: "time_offset",
    __all__[3]: "time_fade",
    __all__[4]: "sound_mix"}

class _augmentation:
    """
    Private base class for different augmentations.
    Prepares general parameters which every augmentation
    needs.

    Parameters
    ----------
    `sample_rate`:
        sample rate of audio file (used for resampling). Default
        is `22050`.
    `function`:
        function which applies the augmentation to the signal
        array. Can be used to reference internal functions or
        functions provided by other modules. Default is `None`.
    `kwargs`:
        Arguments for `function`.

    Notes
    -----
    This is a private base class which is not intened to be
    called on its own. Always call an actual augmentation.
    """

    def __init__(self, sample_rate=22050, function=None, **kwargs):

        self.sample_rate = sample_rate
        self.descriptor = descriptors["_augmentation"]
        self.function = function
        self.kwargs = kwargs
        # use all kwargs values to create a unique ID
        kwargs_vals = []
        for arg in kwargs.values():
            if isinstance(arg, str):
                if "/" in arg or "\\" in arg:
                    arg = self._shortenName(arg)
            kwargs_vals.append(str(arg))
        self.tag = ('_'.join(kwargs_vals)).replace('.', '_')

    def run(self, f_source, f_dest=None, **kwargs):
        """
        Apply the augmentation to a specified file and store
        the result as array or other file.

        Parameters
        ----------
        `f_source`:
            path to file which should be augmented. Must be in wav format.
        `f_dest`:
            path to augmented file which will be saved. If left unspecified,
            a `numpy` array of the augmented signal is returned instead. Pass
            `"auto"` to store in the source with unique names. Default
            is `None`.
        `kwargs`:
            Arguments for executing of `function`.
        """
        self._load(f_source)
        self.f_dest = self._parsePath(f_dest)

        # append kwargs in case some are appended after init
        self.kwargs |= kwargs

        x_new = self._executeFunction()

        if self.f_dest != None:
            sf.write(self.f_dest, x_new, self.sample_rate)
            return None

        else:
            return x_new

    def _load(self, f_source):

        path_details = os.path.splitext(f_source)
        extension = path_details[1]

        if extension != ".wav":
            raise ValueError("File type not supported! Use '.wav' instead.")

        self.f_source = f_source

        x, sr = lr.load(path=self.f_source, sr=self.sample_rate)
        self.signal = x

    def _parsePath(self, f_dest):
        # user wants to save the augmentation
        if f_dest != None:
            if self.f_source == f_dest:
                warnings.warn("Source and save name are the same,\
                    original file will be overwritten!")
                return f_dest
            elif f_dest == "auto":
                return self._autoName(self.descriptor, self.tag)
            else:
                return f_dest

        # user wants a temporary array
        else:
            return None

    def _autoName(self, descriptor, tag):
        if self.f_source == None:
            raise ValueError("No source to autoname from!")
        return self.f_source[:-4] + f"_{descriptor}_{tag}.wav"

    def _executeFunction(self):
        return self.function(self.signal, **self.kwargs)

    def _shortenName(self, f_path):
        return os.path.basename(f_path)[:-4]

class timeStretch(_augmentation):
    """
    Stretch or squeeze a sound while retaining pitch.

    Parameters
    ----------
    `rate`:
        Float. Stretch factor defined by speed of playback. A `rate` < 1
        stretches the sound, a `rate` > 1 squeezes it.
        Default is `1`
    `sample_rate`:
        Int. Sample rate of audio file (used for resampling).
        Default is `22050`
    """

    def __init__(self, rate=1, sample_rate=22050):
        self.rate = rate
        super().__init__(sample_rate=sample_rate,
                         function=lr.effects.time_stretch,
                         rate=rate)
        self.descriptor = descriptors[__all__[0]]

class pitchShift(_augmentation):
    """
    Shift a given input signal's pitch by semitones.

    Parameters
    ----------
    `semitones`:
        Float. Number of western scale semitones to shift up or down.
        Default is `1`
    `sample_rate`:
        Int. Sample rate of audio file (used for resampling).
        Default is `22050`
    """

    def __init__(self, semitones=1, sample_rate=22050):
        self.semitones = semitones
        super().__init__(sample_rate=sample_rate,
                         function=lr.effects.pitch_shift,
                         sr=sample_rate,
                         n_steps=semitones)
        self.descriptor = descriptors[__all__[1]]

class offsetAudio(_augmentation):
    """
    Offset a sound by added dead-time or or by later start.

    Parameters
    ----------
    `s`:
        Float. Offset in seconds. s < 0: skip first samples.
        s > 0: add dead-time to start. Default is `0`.
    `sample_rate`:
        Int. Sample rate of audio file (used for resampling).
        Default is `22050`
    """

    def __init__(self, s=0, sample_rate=22050):
        self.s = s
        super().__init__(sample_rate=sample_rate,
                         function=self._offset,
                         s=s)
        self.descriptor = descriptors[__all__[2]]

    def _offset(self, y, s):

        sample_offs = int(self.sample_rate * abs(s))

        if len(y) <= sample_offs:
            raise ValueError("Offset longer than duration of sound!")
        if s < 0:
            return y[sample_offs:]
        else:
            return np.insert(y, 0, np.zeros(sample_offs))

class fadeAudio(_augmentation):
    """
    Create a logarithmic fade-in or fade-out for a sound.

    Parameters
    ----------
    `s`:
        Float. Fade time in seconds. Default is `0`.
    `direction`:
        String. Direction from which the fade is applied. Default
        is `"in"`.
    `sample_rate`:
        Int. Sample rate of audio file (used for resampling).
        Default is `22050`
    """

    def __init__(self, s=0, direction="in", sample_rate=22050):
        self.s = s
        self.direction = direction
        if direction not in ["in", "out"]:
            raise ValueError(f"parameter '{direction}' not recognized!")

        super().__init__(sample_rate=sample_rate,
                         function=self._fade,
                         s=s,
                         direction=direction)
        self.s = s
        self.direction = direction
        self.descriptor = descriptors[__all__[3]]

    def _fade(self, y, s, direction):

        fade_len = self.sample_rate * s
        x_new = y

        if direction == "out":
            end = len(y)
            start = end - fade_len
            fade_curve = np.logspace(0, -3, fade_len)
            x_new[start:end] *= fade_curve
        else:
            fade_curve = np.logspace(-3, 0, fade_len)
            x_new[0:fade_len] *= fade_curve

        return x_new

class mixAudio(_augmentation):
    """
    Mix two wavefiles at random or specified timestamps.

    Parameters
    ----------
    `ratio`:
        Float. Ratio by which the sounds are mixed.
        `0 <= ratio <= 1`, 1 ignores the noise, 0 the main sound.
        Default is `1`.
    `start_at`:
        Float. Second at which the mixed in sound should be used.
        If unspecified, a random time will be used. Default is `None`.
    `sample_rate`:
        Int. Sample rate of audio file (used for resampling).
        Default is `22050`

    Notes
    -----
    Augmented sound is as long as the original sound of interest,
    not the mixed-in sound
    """

    def __init__(self, ratio=0.5, start_at=None, sample_rate=22050):
        self.ratio = ratio
        self.start_at = start_at
        super().__init__(sample_rate=sample_rate,
                         function=self._mix,
                         ratio=ratio,
                         start_at=start_at)
        self.descriptor = descriptors[__all__[4]]

    def _mix(self, y, f_mix, ratio, start_at):
        y_mix, _ = lr.load(f_mix, sr=self.sample_rate)
        y_len = len(y)
        y_mix_len = len(y_mix)

        # use value of center sample as seed
        if start_at == None:
            rd_value = int(1000*y[int(y_len/2)])
            rd.seed(rd_value)
        else:
            start = int(start_at * self.sample_rate)

        if y_len < y_mix_len:
            if start_at == None:
                start = rd.randint(0, y_mix_len-y_len)
            return y * ratio + y_mix[start:start+y_len] * (1-ratio)
        else:
            if start_at == None:
                start = rd.randint(0, y_len-y_mix_len)
                y_mix_pad = np.append(np.zeros(start), y_mix)
                y_mix_pad = np.append(
                    y_mix_pad, np.zeros(y_len-start-y_mix_len))

            return y * ratio + y_mix_pad * (1-ratio)


import json
from main import db
from main import User, Order, Offer
def open_users():
    sm = []
    with open('data/users.json', 'r', encoding='utf-8') as f:
        data = json.load(f)
        for user in data:
            sm.append(User(id=user['id'],
                           first_name=user['first_name'],
                           last_name=user['last_name'],
                           age=user['age'],
                           email=user['email'],
                           role=user['role'],
                           phone=user['phone']))
    db.session.add_all(sm)
    db.session.commit()

def open_orders():
    sm = []
    with open('data/orders.json', 'r', encoding='utf-8') as f:
        data = json.load(f)
        for order in data:
            sm.append(Order(id=order['id'],
                            name=order['first_name'],
                            description=order['description'],
                            start_date=order['start_date'],
                            end_date=order['end_date'],
                            address=order['address'],
                            price=order['price'],
                            customer_id=order['customer_id'],
                            executor_id=order['executor_id']))
    db.session.add_all(sm)
    db.session.commit()

def open_offers():
    sm = []
    with open('data/offers.json', 'r', encoding='utf-8') as f:
        data = json.load(f)
        for offer in data:
            sm.append(Offer(id=offer['id'], order_id=offer['order_id'], executor_id=offer['executor_id']))
    db.session.add_all(sm)
    db.session.commit()

def all_guides(smt):
    return {
        'id' : smt.id,
        'surname' : smt.surname,
        'full_name' : smt.full_name,
        'tours_count' : smt.tours_count,
        'bio' : smt.bio,
        'is_pro' : smt.is_pro,
        'company' : smt.company
    }


from rest_framework import serializers
from blog.models import Article
from django.contrib.auth import get_user_model
from drf_dynamic_fields import DynamicFieldsMixin

#class AuthorSerializer(serializers.ModelSerializer):     #its about nested serializers( todarto )
#    class Meta:
#        model = get_user_model()
#        fields = ["id","username","last_name","first_name"]

class ArticleSerializer(DynamicFieldsMixin, serializers.ModelSerializer):
#    author = AuthorSerializer()                            #related to above(nested serializers)
#    author = serializers.HyperlinkedIdentityField(view_name='api:authors-detail')   # wrotten for hyperlinks
    def get_author(self, obj):
        return {
             "usename": obj.author.username,
             "first_name": obj.author.first_name,
             "last_name": obj.author.last_name,
        }
    author = serializers.SerializerMethodField("get_author")
    class Meta:
        model = Article
        #fields = ("title", "slug", "author", "content", "publish", "status")
        #exclude = ('created', 'updated')    #means everything except...
        fields = "__all__"
    def validate_title(self, value):
        filter_list = ["javascript","laravel","php"]

        for i in filter_list:
            if i in value:
                raise serializers.ValidationError('plz do not use these words! : {}'.format(i))

class UserSerializer(serializers.ModelSerializer):
    class Meta:
        model = get_user_model()
        fields = "__all__"


import requests
import json
import os

surveyAPI = "" #check email
print("working from where survey, pls enter 1 or 0 as response, 1 beign yes\n\n")
wfh = int(input("Do you prefer complete WFH? "))
if wfh not in [0, 1]:
   print("Bad input %d. Just 1 or 0" %wfh)
   os._exit(-1)

hw = int(input("Do you prefer Hybrid work? "))
if hw not in [0, 1]:
   print("Bad input %d. just 1 or 0" %hw)
   os._exit(-1)

ow = int(input("Do you prefer complete office? "))
if ow not in [0, 1]:
   print("Bad input %d. Just 1 or 0" %ow)
   os._exit(-1)

print("Thanks for participating in survey, uploading your responses")

data = {
  "responses": {
    "Permanent WFH": wfh,
    "Hybrid": hw,
    "Office": ow,
  }
}
response = requests.post(surveyAPI, json.dumps(data))
resp = json.loads(response.content)
print(resp)
survey = resp["body"]
print(survey)


import time

import requests
from selenium import webdriver

# hubstudio_connector.exe --server_mode=http --http_port=6873 --app_id= --group_code= --app_secret=

# 获取webdriver
def startBrowser(hubId):
    url = "http://127.0.0.1:6873/api/v1/browser/start"
    data = {
        "containerCode": hubId
    }
    resp = requests.post(url, json=data).json()
    webdriver_path = resp["data"]["webdriver"]
    port = resp["data"]["debuggingPort"]
    options = webdriver.ChromeOptions()
    options.add_experimental_option("debuggerAddress", '127.0.0.1:' + str(port))
    driver = webdriver.Chrome(webdriver_path, options=options)
    return driver

def stopBrowser(hubId):
    url = "http://127.0.0.1:6873/api/v1/browser/stop?containerCode=" + hubId
    requests.get(url).json()

def open_baidu(driver):
    driver.get("https://www.baidu.com/")
    time.sleep(3)

if __name__ == '__main__':
    # 获取webdriver
    driver = startBrowser('36720345')  # 填写打开环境返回的webdriver和debuggingPort参数
    # 运行脚本
    open_baidu(driver)
    time.sleep(3)
    stopBrowser('36720345')


import torch
from pycocotools.coco import COCO
import fasttext
from cocoTripletDataset import TripletCOCO_Img2Text, TripletCOCO_Text2Img
import os
import torch.nn as nn
from matplotlib import pyplot as plt
from torchvision.models import resnet18, ResNet18_Weights
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights, fasterrcnn_resnet50_fpn_v2
from torch.utils.data import DataLoader
from net import FasterRCNN_Triplet_COCO
from torchvision import transforms
import re
import numpy as np
import wandb

train_jsonPath = r"/ghome/group04/mcv/datasets/COCO/instances_train2014.json"
train_jsonPathCap= r"/ghome/group04/mcv/datasets/COCO/captions_train2014.json"
train_trainImagesPath = r"/ghome/group04/mcv/datasets/COCO/train2014"

val_jsonPath = r"/ghome/group04/mcv/datasets/COCO/instances_val2014.json"
val_jsonPathCap= r"/ghome/group04/mcv/datasets/COCO/captions_val2014.json"
val_trainImagesPath = r"/ghome/group04/mcv/datasets/COCO/val2014"

fasttext_model_path = r"/home/mcv/m5/fasttext_wiki.en.bin"

out_text_model_path = r"best_text_taskb.pth"
out_img_model_path = r"best_img_taskb.pth"

device = "cuda"
batch_size = 32
n_epochs = 4
log_interval = 5
size = (240,320)

text_pre_model = fasttext.load_model(path=fasttext_model_path)
img_pre_model = FasterRCNN_Triplet_COCO().to(device)

# Transform
transformsPretrained = FasterRCNN_ResNet50_FPN_Weights.COCO_V1.transforms()
transformsAp = torch.nn.Sequential(
    transformsPretrained,
    transforms.Resize(size),
)

img_model = nn.Sequential(
    nn.Linear(1024, 256),
    nn.Dropout(),
    nn.ReLU(),
    nn.Linear(256, 128)
).to(device).eval()

text_model = nn.Sequential(
    nn.Linear(300, 256),
    nn.Dropout(),
    nn.ReLU(),
    nn.Linear(256, 128)
).to(device).eval()

train_dataset = TripletCOCO_Text2Img(train_trainImagesPath, os.listdir(train_trainImagesPath), train_jsonPath, train_jsonPathCap, transformsAp)
val_dataset = TripletCOCO_Text2Img(val_trainImagesPath, os.listdir(val_trainImagesPath), val_jsonPath, val_jsonPathCap, transformsAp)

def sentence2vector(sentence, model):
    sentence = re.sub(r'[^a-zA-Z]', '', sentence).lower()
    words = sentence.split()
    vectors = [model.get_word_vector(word) for word in words]
    summary_vector = np.mean(vectors, axis=0)

    return summary_vector

criterion = nn.TripletMarginLoss()
optimizer = torch.optim.Adam(
    list(img_model.parameters()) + list(text_model.parameters()),
    lr=3e-4
)

def prep_sentence(sentence, text_pre_model):
    sentence = text_pre_model.get_sentence_vector(
        re.sub(r'[^a-zA-Z ]', '', sentence.lower())
    )
    return sentence

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

wandb.init(project="text2img")
min_val_loss = 100
for epoch in range(n_epochs):
    img_model.train()
    text_model.train()

    for img, pos, neg in train_loader:

        img = img.to(device)
        pos = [prep_sentence(p, text_pre_model) for p in pos]
        neg = neg.to(device)
        pos = torch.tensor(pos, device=device)
        with torch.no_grad():
            img_feature_vector = img_pre_model(img)
            neg_feature_vector = img_pre_model(neg)

        img_common = img_model(img_feature_vector.reshape((-1, img_feature_vector.shape[1])))
        neg_img = img_model(neg_feature_vector.reshape((-1, neg_feature_vector.shape[1])))
        pos_common  = text_model(pos)

        loss = criterion(pos_common,img_common, neg_img)
        img_model.zero_grad()
        text_model.zero_grad()
        loss.backward()
        optimizer.step()

        #if epoch % log_interval == 0:
        print(f"Train loss: {loss.item()}")
        wandb.log({"train_loss": loss})

    img_model.eval()
    text_model.eval()
    losses = []
    for img, pos, neg in val_loader:
        img = img.to(device)
        pos = [prep_sentence(p, text_pre_model) for p in pos]
        neg = neg.to(device)
        pos = torch.tensor(pos, device=device)
        with torch.no_grad():
            img_feature_vector = img_pre_model(img)
            neg_feature_vector = img_pre_model(neg)

        img_common = img_model(img_feature_vector.reshape((-1, img_feature_vector.shape[1])))
        neg_img = img_model(neg_feature_vector.reshape((-1, neg_feature_vector.shape[1])))
        pos_common  = text_model(pos)

        loss = criterion(pos_common,img_common, neg_img)
        losses.append(loss.item())
    val_loss = torch.mean(torch.tensor(losses))
    if val_loss < min_val_loss:
        min_val_loss = val_loss
        torch.save(img_model.state_dict(), out_img_model_path)
        torch.save(text_model.state_dict(), out_text_model_path)
    print(f"Val loss: {val_loss}")
    wandb.log({"val_loss": val_loss})