#!/usr/bin/env python3

import os
import unittest
from unittest import mock

from quibble import cmd

class CmdTest(unittest.TestCase):

    @mock.patch.dict('os.environ', clear=True)
    def test_projects_to_clone(self):
        q = cmd.QuibbleCmd()
        self.assertEqual(
            q.set_repos_to_clone(),
            ['mediawiki/core', 'mediawiki/skins/Vector'],
            'Incorrect repos to clone')

    @mock.patch.dict('os.environ', clear=True)
    def test_projects_to_clone_with_vendor(self):
        q = cmd.QuibbleCmd()
        self.assertEqual(
            q.set_repos_to_clone(clone_vendor=True),
            ['mediawiki/core', 'mediawiki/skins/Vector', 'mediawiki/vendor'],
            'Incorrect repos to clone')

    @mock.patch.dict('os.environ', clear=True)
    def test_projects_to_clone_appends_projects(self):
        q = cmd.QuibbleCmd()
        self.assertEqual(
            q.set_repos_to_clone(projects=[
                'mediawiki/extensions/BoilerPlate',
                'mediawiki/extensions/Example',
                ]),
            ['mediawiki/core', 'mediawiki/skins/Vector',
             'mediawiki/extensions/BoilerPlate',
             'mediawiki/extensions/Example'])

    def test_set_repos_to_clone_with_env(self):
        env = {
            'SKIN_DEPENDENCIES': 'mediawiki/skins/Monobook',
            'EXT_DEPENDENCIES':
                'mediawiki/extensions/One\\nmediawiki/extensions/Two',
        }
        with mock.patch.dict('os.environ', env, clear=True):
            q = cmd.QuibbleCmd()
            self.assertEqual([
                'mediawiki/core',  # must be first
                'mediawiki/skins/Vector',
                'mediawiki/skins/Monobook',
                'mediawiki/extensions/One',
                'mediawiki/extensions/Two',
                ], q.set_repos_to_clone())

    def test_set_repos_to_clone_adds_ZUUL_PROJECT(self):
        env = {'ZUUL_PROJECT': 'mediawiki/extensions/ZuulProjectEnvVar'}
        with mock.patch.dict('os.environ', env, clear=True):
            q = cmd.QuibbleCmd()
            self.assertEqual([
                'mediawiki/core',  # must be first
                'mediawiki/skins/Vector',
                'mediawiki/extensions/ZuulProjectEnvVar',
                ], q.set_repos_to_clone())

    def test_set_repos_to_clone_does_not_duplicate_hardcoded_repos(self):
        hardcoded_repos = [
            'mediawiki/core',
            'mediawiki/skins/Vector',
            ]

        for repo in hardcoded_repos:
            with mock.patch.dict('os.environ', {'ZUUL_PROJECT': repo},
                                 clear=True):
                q = cmd.QuibbleCmd()
                self.assertEqual([
                    'mediawiki/core',  # must be first
                    'mediawiki/skins/Vector',
                    ], q.set_repos_to_clone())

    @mock.patch('quibble.is_in_docker', return_value=False)
    def test_args_defaults(self, _):
        args = cmd.QuibbleCmd().parse_arguments([])

        self.assertEqual('ref', args.git_cache)
        self.assertEqual(os.getcwd(), args.workspace)
        self.assertEqual(os.path.join(os.getcwd(), 'log'), args.log_dir)

    @mock.patch('quibble.is_in_docker', return_value=True)
    def test_args_defaults_in_docker(self, _):
        args = cmd.QuibbleCmd().parse_arguments([])

        self.assertEqual('/srv/git', args.git_cache)
        self.assertEqual('/workspace', args.workspace)
        self.assertEqual('/workspace/log', args.log_dir)

    @mock.patch.dict(os.environ, clear=True)
    def test_setup_environment(self):
        q = cmd.QuibbleCmd()
        q.workspace = '/testworkspace'
        q.mw_install_path = ''
        q.log_dir = ''

        with mock.patch('quibble.is_in_docker', return_value=True):
            # In Docker we always use self.workspace
            q.setup_environment()
            self.assertEqual(os.environ['WORKSPACE'], '/testworkspace')
            with mock.patch.dict(os.environ, {'WORKSPACE': '/fromenv'},
                                 clear=True):
                # In Docker, ignore $WORKSPACE
                q.setup_environment()
                self.assertEqual(os.environ['WORKSPACE'], '/testworkspace')

        with mock.patch('quibble.is_in_docker', return_value=False):
            q.setup_environment()
            self.assertEqual(os.environ['WORKSPACE'], '/testworkspace')

            with mock.patch.dict(os.environ, {'WORKSPACE': '/fromenv'},
                                 clear=True):
                # When not in Docker, we honor $WORKSPACE
                q.setup_environment()
                self.assertEqual(os.environ['WORKSPACE'], '/fromenv')

    @mock.patch.dict(os.environ, clear=True)
    def test_setup_environment_has_log_directories(self):
        q = cmd.QuibbleCmd()
        q.workspace = '/workspace'
        q.mw_install_path = ''
        q.log_dir = '/mylog'

        q.setup_environment()

        self.assertIn('LOG_DIR', os.environ)
        self.assertIn('MW_LOG_DIR', os.environ)
        self.assertEqual(os.environ['LOG_DIR'], '/mylog')
        self.assertEqual(os.environ['MW_LOG_DIR'], '/mylog')

    def test_isCoreOrVendor(self):
        q = cmd.QuibbleCmd()
        self.assertTrue(q.isCoreOrVendor('mediawiki/core'))
        self.assertTrue(q.isCoreOrVendor('mediawiki/vendor'))
        self.assertFalse(q.isCoreOrVendor('mediawiki/extensions/Foo'))
        self.assertFalse(q.isCoreOrVendor('mediawiki/skins/Bar'))

    def test_isExtOrSkin(self):
        q = cmd.QuibbleCmd()
        q.isExtOrSkin
        self.assertTrue(q.isExtOrSkin('mediawiki/extensions/Foo'))
        self.assertTrue(q.isExtOrSkin('mediawiki/skins/Bar'))
        self.assertFalse(q.isExtOrSkin('mediawiki/core'))
        self.assertFalse(q.isExtOrSkin('mediawiki/vendor'))

    def test_should_run_accepts_all_stages_by_default(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=[])
        self.assertTrue(
            all(map(q.should_run, q.stages)),
            'must runs all stages by default')

    def test_should_run_runall_accepts_all_stages(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=['--run=all'])
        self.assertTrue(
            all(map(q.should_run, q.stages)),
            '--run=all runs all stages')

    def test_should_run_skippall_runs_no_stage(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=['--skip=all'])
        self.assertFalse(
            any(map(q.should_run, q.stages)),
            '--skip=all skips all stages')

    def test_should_run_skips_a_stage(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=['--skip=phpunit'])
        self.assertFalse(
            q.should_run('phpunit'),
            '--skip skips the stage')
        stages_to_run = [s for s in q.stages
                         if s != 'phpunit']
        self.assertTrue(
            all(map(q.should_run, stages_to_run)),
            'Must runs all non skipped stages')

    def test_should_run_running_a_single_stage(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=['--run=phpunit'])
        self.assertTrue(
            q.should_run('phpunit'),
            '--run runs the stage')
        stages_to_skip = [s for s in q.stages
                          if s != 'phpunit']
        self.assertFalse(
            any(map(q.should_run, stages_to_skip)),
            'Must not run any other stages')

    def test_commands_skip_all_stages(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=['--commands=/bin/true'])
        self.assertFalse(
            any(map(q.should_run, q.stages)),
            'User commands must skip all stages')

    def test_project_branch_arg(self):
        q = cmd.QuibbleCmd()
        q.args = q.parse_arguments(args=[])
        self.assertEquals([], q.args.project_branch)


Counter = 0

def when_started():
    global Counter
    # The counter variable will help us keep track of how many times we have gone through the loop.
    Counter = 0
    for repeat_count in range(10):
        Counter = Counter + 1
        drivetrain.turn_for(RIGHT, 90, DEGREES)
        pen.move(DOWN)
        drivetrain.drive_for(FORWARD, 1800, MM)
        pen.move(UP)
        # Once we have looped 10 times, we will prevent the robot from trying to move up to another line.
        if Counter < 10:
            drivetrain.drive_for(REVERSE, 1800, MM)
            drivetrain.turn_for(LEFT, 90, DEGREES)
            drivetrain.drive_for(FORWARD, 200, MM)
        wait(5, MSEC)

vr_thread(when_started())


from itertools import chain

from odoo import api, fields, models, tools, _
from odoo.exceptions import UserError, ValidationError
from odoo.tools.float_utils import float_compare, float_is_zero, float_round

class Pricelist(models.Model):
    _inherit = "product.pricelist"

    pricelist_move_id = fields.Char(string='明细行选择')

# 给订单明细行选择价格时添加计量单位约束
    @api.multi
    def _compute_price_rule(self, products_qty_partner, date=False, uom_id=False):
        """ Low-level method - Mono pricelist, multi products
        Returns: dict{product_id: (price, suitable_rule) for the given pricelist}

        If date in context: Date of the pricelist (%Y-%m-%d)

            :param products_qty_partner: list of typles products, quantity, partner
            :param datetime date: validity date
            :param ID uom_id: intermediate unit of measure
        """
        # print(products_qty_partner)
        self.ensure_one()
        if not date:
            date = self._context.get('date') or fields.Date.context_today(self)
        if not uom_id and self._context.get('uom'):
            uom_id = self._context['uom']
        if uom_id:
            # rebrowse with uom if given
            products = [item[0].with_context(uom=uom_id) for item in products_qty_partner]

            products_qty_partner = [(products[index], data_struct[1], data_struct[2]) for index, data_struct in
                                    enumerate(products_qty_partner)]

        else:
            products = [item[0] for item in products_qty_partner]

        if not products:
            return {}

        categ_ids = {}
        for p in products:
            categ = p.categ_id
            while categ:
                categ_ids[categ.id] = True
                categ = categ.parent_id
        categ_ids = list(categ_ids)

        is_product_template = products[0]._name == "product.template"
        if is_product_template:
            prod_tmpl_ids = [tmpl.id for tmpl in products]
            # all variants of all products
            prod_ids = [p.id for p in
                        list(chain.from_iterable([t.product_variant_ids for t in products]))]
        else:
            prod_ids = [product.id for product in products]
            prod_tmpl_ids = [product.product_tmpl_id.id for product in products]

        # Load all rules
        self._cr.execute(
            'SELECT item.id '
            'FROM product_pricelist_item AS item '
            'LEFT JOIN product_category AS categ '
            'ON item.categ_id = categ.id '
            'WHERE (item.product_tmpl_id IS NULL OR item.product_tmpl_id = any(%s))'
            'AND (item.product_id IS NULL OR item.product_id = any(%s))'
            'AND (item.categ_id IS NULL OR item.categ_id = any(%s)) '
            'AND (item.pricelist_id = %s) '
            'AND (item.date_start IS NULL OR item.date_start<=%s) '
            'AND (item.date_end IS NULL OR item.date_end>=%s)'
            'ORDER BY item.applied_on, item.min_quantity desc, categ.complete_name desc, item.id desc',
            (prod_tmpl_ids, prod_ids, categ_ids, self.id, date, date))
        # NOTE: if you change `order by` on that query, make sure it matches
        # _order from model to avoid inconstencies and undeterministic issues.

        sale_uoms = self._context.get('sale_uom')

        item_ids = [x[0] for x in self._cr.fetchall()]
        # print(item_ids,'item_ids')
        items = self.env['product.pricelist.item'].browse(item_ids)
        # print(items, 'items')
        results = {}
        for product, qty, partner in products_qty_partner:
            # print(len(products_qty_partner),6354)
            results[product.id] = 0.0
            suitable_rule = False

            # Final unit price is computed according to `qty` in the `qty_uom_id` UoM.
            # An intermediary unit price may be computed according to a different UoM, in
            # which case the price_uom_id contains that UoM.
            # The final price will be converted to match `qty_uom_id`.
            qty_uom_id = self._context.get('uom') or product.uom_id.id
            price_uom_id = product.uom_id.id
            qty_in_product_uom = qty
            if qty_uom_id != product.uom_id.id:
                try:
                    qty_in_product_uom = self.env['uom.uom'].browse([self._context['uom']])._compute_quantity(qty,
                                                                                product.uom_id)
                except UserError:
                    # Ignored - incompatible UoM in context, use default product UoM
                    pass

            # if Public user try to access standard price from website sale, need to call price_compute.
            # TDE SURPRISE: product can actually be a template
            price = product.price_compute('list_price')[product.id]

            price_uom = self.env['uom.uom'].browse([qty_uom_id])

            # sale_context = self.env['sale.order.line'].context_get()
            # print(sale_context,'sale_context')

            for rule in items:
                # print(rule.product_sale_uom,'rl')
                # print(sale_uoms, 'sale_uoms')
                if rule.product_measuring_unit_name == sale_uoms:
                    # print(rule.product_sale_uom, '单位相等')
                    if rule.min_quantity and qty_in_product_uom < rule.min_quantity:
                        continue
                    if is_product_template:
                        if rule.product_tmpl_id and product.id != rule.product_tmpl_id.id:
                            continue
                        if rule.product_id and not (
                                product.product_variant_count == 1 and product.product_variant_id.id == rule.product_id.id):
                            # product rule acceptable on template if has only one variant
                            continue
                    else:
                        if rule.product_tmpl_id and product.product_tmpl_id.id != rule.product_tmpl_id.id:
                            continue
                        if rule.product_id and product.id != rule.product_id.id:
                            continue
                if rule.product_measuring_unit_name == sale_uoms:
                    if rule.categ_id:
                        cat = product.categ_id
                        while cat:
                            if cat.id == rule.categ_id.id:
                                break
                            cat = cat.parent_id
                        if not cat:
                            continue
                if rule.product_measuring_unit_name == sale_uoms:
                    if rule.base == 'pricelist' and rule.base_pricelist_id:
                        price_tmp = rule.base_pricelist_id._compute_price_rule([(product, qty, partner)])[product.id][
                            0]  # TDE: 0 = price, 1 = rule
                        price = rule.base_pricelist_id.currency_id._convert(price_tmp, self.currency_id,
                                                                            self.env.user.company_id, date, round=False)
                    else:
                        # if base option is public price take sale price else cost price of product
                        # price_compute returns the price in the context UoM, i.e. qty_uom_id
                        price = product.price_compute(rule.base)[product.id]

                convert_to_price_uom = (lambda price: product.uom_id._compute_price(price, price_uom))
                # print(convert_to_price_uom,'convert_to_price_uom')

                if rule.product_measuring_unit_name == sale_uoms:
                    if price is not False:
                        if rule.compute_price == 'fixed':
                            price = convert_to_price_uom(rule.fixed_price)
                        elif rule.compute_price == 'percentage':
                            price = (price - (price * (rule.percent_price / 100))) or 0.0
                        else:
                            # complete formula
                            price_limit = price
                            price = (price - (price * (rule.price_discount / 100))) or 0.0
                            if rule.price_round:
                                price = tools.float_round(price, precision_rounding=rule.price_round)

                            if rule.price_surcharge:
                                price_surcharge = convert_to_price_uom(rule.price_surcharge)
                                price += price_surcharge

                            if rule.price_min_margin:
                                price_min_margin = convert_to_price_uom(rule.price_min_margin)
                                price = max(price, price_limit + price_min_margin)

                            if rule.price_max_margin:
                                price_max_margin = convert_to_price_uom(rule.price_max_margin)
                                price = min(price, price_limit + price_max_margin)
                        suitable_rule = rule
                    break
            # Final price conversion into pricelist currency
            if suitable_rule and suitable_rule.compute_price != 'fixed' and suitable_rule.base != 'pricelist':
                price = product.currency_id._convert(price, self.currency_id, self.env.user.company_id, date,
                                                     round=False)
            results[product.id] = (price, suitable_rule and suitable_rule.id or False)

        return results

class StockingPicking(models.Model):
    _inherit = 'stock.picking'

# 判断所输入完成数量与所选择的计量单位是否符合
    @api.multi
    def button_validate(self):
        self.ensure_one()
        if not self.move_lines and not self.move_line_ids:
            raise UserError(_('Please add some items to move.'))

        # If no lots when needed, raise error
        picking_type = self.picking_type_id
        precision_digits = self.env['decimal.precision'].precision_get('Product Unit of Measure')
        no_quantities_done = all(
            float_is_zero(move_line.qty_done, precision_digits=precision_digits) for move_line in self.move_line_ids)
        no_reserved_quantities = all(
            float_is_zero(move_line.product_qty, precision_rounding=move_line.product_uom_id.rounding) for move_line in
            self.move_line_ids)
        if no_reserved_quantities and no_quantities_done:
            raise UserError(_(
                'You cannot validate a transfer if no quantites are reserved nor done. To force the transfer, switch in edit more and encode the done quantities.'))

        for line in self.move_ids_without_package:
            if line.quantity_done == 0:
                break
            print(line.quantity_done)
            # uom_env = self.env['uom.uom'].search([('name', '=', line.product_unit)])
            if line.product_unit.factor_inv == 0:
                break

            result = line.quantity_done / line.product_unit.factor_inv
            print(result, 'result')
            result1 = int(result)
            if result != result1:
                raise ValidationError('您输入的完成数量与所选择的销售单位不符')

        if picking_type.use_create_lots or picking_type.use_existing_lots:
            lines_to_check = self.move_line_ids
            if not no_quantities_done:
                lines_to_check = lines_to_check.filtered(
                    lambda line: float_compare(line.qty_done, 0,
                                               precision_rounding=line.product_uom_id.rounding)
                )

            for line in lines_to_check:
                product = line.product_id
                if product and product.tracking != 'none':
                    if not line.lot_name and not line.lot_id:
                        raise UserError(
                            _('You need to supply a Lot/Serial number for product %s.') % product.display_name)

        if no_quantities_done:
            view = self.env.ref('stock.view_immediate_transfer')
            wiz = self.env['stock.immediate.transfer'].create({'pick_ids': [(4, self.id)]})
            return {
                'name': _('Immediate Transfer?'),
                'type': 'ir.actions.act_window',
                'view_type': 'form',
                'view_mode': 'form',
                'res_model': 'stock.immediate.transfer',
                'views': [(view.id, 'form')],
                'view_id': view.id,
                'target': 'new',
                'res_id': wiz.id,
                'context': self.env.context,
            }

        if self._get_overprocessed_stock_moves() and not self._context.get('skip_overprocessed_check'):
            view = self.env.ref('stock.view_overprocessed_transfer')
            wiz = self.env['stock.overprocessed.transfer'].create({'picking_id': self.id})
            return {
                'type': 'ir.actions.act_window',
                'view_type': 'form',
                'view_mode': 'form',
                'res_model': 'stock.overprocessed.transfer',
                'views': [(view.id, 'form')],
                'view_id': view.id,
                'target': 'new',
                'res_id': wiz.id,
                'context': self.env.context,
            }

        # Check backorder should check for other barcodes
        if self._check_backorder():
            return self.action_generate_backorder_wizard()
        self.action_done()
        return


import pandas as pd
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import pyplot as plt
from DBConnections import databases_connection
from datetime import datetime, timedelta
from pandas.tools.plotting import table

class SpreadQuantiles(object):
    def __init__(self):
        self.vertica_cursor = databases_connection.vertica_cursor
        self.instrument = 'USDRUB_TOM'
        self.clean_prices_map = {'0.5 mio': ['119:', 0.000], '1 mio': ['90:', 0.001], '3 mio': ['91:', 0.002],
                                 '5 mio': ['92:', 0.003], '10 mio': ['120:', 0.004]}
        self.default_path = '/home/glebyarnykh/db/PreliminaryAnalysis/SpreadQuantiles8.pdf'

    def get_trading_days(self, start_date, end_date):
        pass

    def get_daily_spreads_quantiles(self, date):
        final_pdf = PdfPages(self.default_path)
        params = {'start': date.replace(hour=10, minute = 30), 'end': date.replace(hour=16, minute=30), 'instr': self.instrument}
        for size in self.clean_prices_map.keys():
            price_group = self.clean_prices_map[size][0]
            avg_margin = self.clean_prices_map[size][1]
            params['group'] = price_group
            query = "SELECT Bid, Ask FROM prices where instrument = :instr and time>:start" \
                    " and time<:end and PricingGroup = :group;"
            self.vertica_cursor.execute(query, params)
            data = pd.DataFrame(self.vertica_cursor.fetchall(), columns=['Bid', 'Ask'])
            data.loc[:, 'Spread'] = (data.loc[:, 'Ask'] - data.loc[:, 'Bid']) + avg_margin*2
            fig, ax = plt.subplots(1, 1)
            data.loc[:, 'Spread'].hist(ax = ax, bins = 30, normed=True)
            table(ax, data.loc[:, 'Spread'].describe().round(3), loc='upper left', colWidths=[0.2])
            fig.suptitle(size, fontsize=15)
            final_pdf.savefig(fig)
            fig.clf()
        final_pdf.close()

if __name__ == '__main__':
    check = SpreadQuantiles()
    check.get_daily_spreads_quantiles(datetime(2016, 7, 8))


#!/usr/bin/env python
'''
A quick and dirty script to visualize results from our example fortran
program.  Opens the given file, reads data into a 2D array and time/space
grid into 1D arrays, and creates a contour plot to illustrate the result.
'''

from argparse import ArgumentParser
parser = ArgumentParser(description=__doc__)
# Add arguments:
parser.add_argument('filename', help="Name of file to open and plot.",
                    type=str)
args = parser.parse_args()

# The usual imports:
import re
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.ticker import LogLocator, LogFormatterMathtext

# Open file.
f=open(args.filename, 'r')

# Parse header.
title=f.readline().strip() # Use header line as plot title.

# Use regular expressions to parse the rest of the header.
# Regexs made this step a lot easier!
# The domain as a numpy array:
domain = np.array(re.findall('(\d+\.\d*)', f.readline()), dtype=float)
# size of domain in space and time (as integers):
nx, nt = re.search('(\d+)x(\d+)', f.readline()).groups()  # number of points.
nx, nt = int(nx), int(nt)

# Create containers for x-grid, time-grid, and 2D results array:
results = np.zeros( (nx, nt) )
t = np.zeros(nt)

# Read spatial grid (just on the first line, skip sub-header).
# Convert into floating point integer numpy array.
x = np.array(f.readline().split()[1:], dtype=float)

# Read rest of file:
for i, l in enumerate(f.readlines()):
    parts = l.split()
    t[i] = parts[0]
    results[:,i] = parts[1:]

# Some ranges/values for the plot.  We use these to control the colorbar.
minlog = 1E-3
maxlog = 1.0
lev_exp = np.arange(np.log10(minlog), np.log10(maxlog), 0.1)
levs = np.power(10, lev_exp)

# Create a figure and a plot:
fig=plt.figure()
ax =fig.add_subplot(111)

# Sometimes, we get 0 or below 0 results, which doesn't work with
# log axes (or log-scale contours).  Tricky indexing!!
results[results<minlog] = minlog

# Create a filled contour.  Note how we use the levels and limits from
# above to carefully adjust and control our plot.  See the Matplotlib docs!
cont =ax.contourf(t, x, results, levs,
                  norm=LogNorm(), cmap='hot')
# Create a colorbar for the plot, use math text ticks.
cbar=fig.colorbar(cont, ax=ax, ticks=LogLocator(),
                  format=LogFormatterMathtext())
cbar.set_label('Intensity (arbitrary units)')

# Finally, label axes/plot.
ax.set_ylabel('X (arbitrary units)')
ax.set_xlabel('Time ($s$)')
ax.set_title(title)

# These lines are useful for updating the figure:
if plt.isinteractive():
    plt.draw()
else:
    plt.show()


import sys, os, time, datetime, curses, logging
sys.path.insert(0, os.path.abspath(__file__+'/../..'))
import argparse, glob, re, cv2
from tqdm import tqdm
import numpy as np
import pynvml, subprocess
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')
import torch
from torch import nn
import torch.nn.functional as F

from maskpyramid.config import cfg
from maskpyramid.backbone.unet_baseline.unet_purebase import UNet, down, up
# from maskpyramid.modeling.backbone import resnet
from dataloaders.build import make_data_loader
from maskpyramid.utils.metrics import Evaluator
from maskpyramid.utils.lr_scheduler import LR_Scheduler
from maskpyramid.utils.tools import setup_logger, mkdir, collect_env_info
from maskpyramid.utils.tools import TensorboardSummary, flatten_list

class KeepDown(nn.Module):
    def __init__(self, cfg):
        super(KeepDown, self).__init__()
        self.down5 = down(512, 512)
        self.down6 = down(512, 512)
        self.down7 = down(512, 512)
        self.down8 = down(512, 512)
        self.down9 = down(512, 512)

    def forward(self, x5):
        x6 = self.down5(x5)
        x7 = self.down6(x6)
        x8 = self.down7(x7)
        x9 = self.down8(x8)

        return [x6, x7, x8, x9]

class CategoryPyramid(nn.Module):
    def __init__(self, cfg):
        super(CategoryPyramid, self).__init__()
        num_classes = cfg.DATALOADER.NUM_CLASSES
        self.f_down_1x1 = nn.ModuleList([nn.ModuleList([nn.Conv2d(512,256,1) for \
            i in range(j+1)]) for j in range(6)])
        self.f_down_1x1_129 = nn.Conv2d(256,256,1)
        self.up1 = up(512, 256)    # 3
        self.up2 = up(512, 256)    # 5
        self.up3 = up(512, 256)    # 9
        self.up4 = up(512, 256)    # 17
        self.up5 = up(512, 256)    # 33
        self.up6 = up(512, 256)    # 3
        self.up129 = up(512, 256)    # 3
        self.out_conv = nn.Conv2d(256, 1, 1)

        self.up_list = [None, self.up1, self.up2, self.up3, self.up4, self.up5, self.up6]
        # self.up_list = nn.ModuleList([None, self.up1, self.up2, self.up3,
        #     self.up4, self.up5, self.up6])

class MaskPyramids(nn.Module):
    def __init__(self, cfg):
        super(MaskPyramids, self).__init__()
        self.cfg = cfg
        num_classes = cfg.DATALOADER.NUM_CLASSES
        self.sematic_bone = UNet(num_classes)
        self.sematic_bone.load_state_dict(torch.load('maskpyramid/backbone/unet_baseline/checkpoint_MIOU54.4.pth.tar')['state_dict'])

        self.keep_down = KeepDown(cfg)
        cat_conv = [None for _ in range(11)] + [CategoryPyramid(cfg) for _ in range(8)]
        self.cat_conv = nn.ModuleList(cat_conv)
        self.sematic_criteron = nn.CrossEntropyLoss(ignore_index=cfg.DATALOADER.IGNORE_INDEX)
        self.sem_ins_criteron = nn.CrossEntropyLoss(ignore_index=cfg.DATALOADER.IGNORE_INDEX)

    def prepare_targets(self, target, down_size, mode='sematic', pyramids=[]):

        N, img_h, img_w = target['label'].shape
        h_ratio = int(np.ceil(img_h / down_size[0]))
        w_ratio = int(np.ceil(img_w / down_size[1]))
        h_match_size = h_ratio * down_size[0]
        w_match_size = w_ratio * down_size[1]

        inst_mode_num_stat = F.interpolate(target['instance'].unsqueeze(1).float(), (h_match_size,
            w_match_size), mode='nearest').squeeze(1).long().view((N, down_size[0],
            h_ratio, down_size[1], w_ratio)).permute(0,1,3,2,4).contiguous().view((N,
            )+down_size+(-1,))
        mode_num_instance = inst_mode_num_stat.mode(-1)[0]

        if mode == 'instance':
            return mode_num_instance

        elif mode == 'sematic':
            mode_num_sematic = mode_num_instance.clone()
            for bid in range(len(mode_num_sematic)):
                for inst_num in mode_num_instance[bid].unique():
                    mode_num_sematic[bid][mode_num_instance[bid] == inst_num] = target['label']\
                        [bid][target['instance'][bid] == inst_num].unique().item()
            return mode_num_sematic
        else:
            import pdb; pdb.set_trace()
        return 0

    def compute_mask(self, level, bid, feature, pyramids):
        masks = []
        involved_features = feature[::-1][level:]
        for pyr in pyramids:
            cat_conv = self.cat_conv[pyr.tar_cat]
            feat_cover_mask, p_rge = pyr.get_cover_mask(level)
            f_down_0 = involved_features[0][[bid], :, p_rge[0]:p_rge[1], p_rge[2]:p_rge[3]]
            x = cat_conv.f_down_1x1[level][0](f_down_0)
            # 0-3 1-5 2-9 3-17 4-33 5-65 6-129
            for i in range(1, 6-level):
                feat_cover_mask, p_rge = pyr.get_cover_mask(level+i)
                f_down = involved_features[i][[bid], :, p_rge[0]:p_rge[1], p_rge[2]:p_rge[3]]
                x_down = cat_conv.f_down_1x1[level+i][i](f_down)
                x = cat_conv.up_list[i](x, x_down)
            feat_cover_mask, p_rge = pyr.get_cover_mask(6)
            f_down_129 = feature[2][[bid], :, p_rge[0]:p_rge[1], p_rge[2]:p_rge[3]]
            x_down_129 = cat_conv.f_down_1x1_129(f_down_129)
            x_up_129 = cat_conv.up129(x, x_down_129)
            pad_129 = feature[0].new_zeros((1,x_up_129.shape[1], 129, 129))
            pad_129[0,:,p_rge[0]:p_rge[1], p_rge[2]:p_rge[3]] = x_up_129
            mask_129 = cat_conv.out_conv(pad_129)
            pyr.set_mask_129(mask_129)
            masks.append(mask_129)

        return masks

    def rend_on_image_v8(self, image_np, masks_np, labels_list):
        # label 从0开始，0 表示unlabeled
        alpha = 0.5
        color_bias = 1
        # import pdb; pdb.set_trace()
        # print('masks_np unique',np.unique(masks_np), 'labels_list', labels_list)
        class_count = [0 for _ in range(20)]
        colors = np.array([
            # [0, 100, 0],
            [128, 64, 128],
            [244, 35, 232],
            [70, 70, 70],
            [102, 102, 156],
            [190, 153, 153],
            [153, 153, 153],
            [250, 170, 30],
            [220, 220, 0],
            [107, 142, 35],
            [152, 251, 152],
            [0, 130, 180],
            [220, 20, 60],
            [255, 0, 0],
            [0, 0, 142],
            [0, 0, 70],
            [0, 60, 100],
            [0, 80, 100],
            [0, 0, 230],
            [119, 11, 32]])
        masked = image_np.copy()
        for i, label in enumerate(labels_list):
            if label == -1:
                continue
            mask_idx = np.nonzero(masks_np == i)
            masked[mask_idx[0], mask_idx[1], :] *= 1.0 - alpha
            color = (colors[label]+class_count[label]*color_bias)/255.0

            masked[mask_idx[0], mask_idx[1], :] += alpha * color
            contours, hierarchy = cv2.findContours(
                (masks_np == i).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
            )
            masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), 1)
            class_count[label] += 1

        return masked

    def show_image_v8(self, images, output, targets, tool='cv2'):
        N,C,H,W = images.shape
        masked_imgs = []
        origin_image_ON = True
        ori_sematic_target_ON = False
        ori_instance_target_ON = True
        car_only_ON = False
        pred_sematic_ON = True
        pred_instance_ON = False
        pred_stat_ON = True
        for i, image in enumerate(images):
            parts_to_show = []
            images_np = image.permute(1,2,0).detach().cpu().numpy()
            images_np = ((images_np - images_np.min()) / (images_np.max() - images_np.min()))[...,::-1]
            if origin_image_ON:
                parts_to_show.append(images_np)
            semat_a_target, insts_a_target = targets['label'][i], targets['instance'][i]

            # targets prepared in forwards
            # out_one = output['in_dict_watches'][i]
            # target_fwd_label = (out_one['target_levels']['labels']+1).tolist()
            sematic_target = output['targets']['label'][i]
            instacne_target = output['targets']['instance'][i]
            # sematic_out_513 = output['sematic_out'][6]
            sematic_out_513 = output['sematic_out']

            if ori_sematic_target_ON:
                sematic_target_np = sematic_target.detach().cpu().numpy()
                ori_sematic_target = self.rend_on_image_v8(images_np, sematic_target_np, range(19))
                parts_to_show.append(ori_sematic_target)

            if ori_instance_target_ON:
                instacne_target_np = instacne_target.detach().cpu().numpy()
                ins_cats = [sematic_target[instacne_target == ins_idx].unique().item() for ins_idx in instacne_target.unique()]
                if car_only_ON:
                    ins_cats = [13 if label==13 else -1 for label in ins_cats]
                ori_instacne_target = self.rend_on_image_v8(images_np, instacne_target_np, ins_cats)
                parts_to_show.append(ori_instacne_target)

            # import pdb; pdb.set_trace()
            # prediction    - sematic
            if pred_sematic_ON:
                sematic_predict_np = sematic_out_513[i].max(0)[1].detach().cpu().numpy()
                sematic_predict = self.rend_on_image_v8(images_np, sematic_predict_np, range(19))
                parts_to_show.append(sematic_predict)

            if not self.cfg.SOLVER.SEMATIC_ONLY:
                parymids = output['ins_pyramids']
                # prediction    - instance
                if pred_instance_ON:
                    # import pdb; pdb.set_trace()
                    instance_predict = torch.cat([pyr.get_mask(513) for pyr in  parymids[i]], dim=1) if len(parymids[i]) else torch.empty((1,1,513,513))
                    instance_predict_513 = F.interpolate(instance_predict, 513, mode='bilinear', align_corners=True)
                    instance_predict_np = instance_predict_513[0].max(0)[1].detach().cpu().numpy()
                    ins_cats_predict = [pyr.tar_cat for pyr in  parymids[i]]
                    # import pdb; pdb.set_trace()
                    pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, ins_cats_predict)
                    # pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, range(19))
                    if pred_stat_ON:
                        cv2.putText(pred_instance,'#:{}'.format(len(parymids[i])),(30,30),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1)
                        for pyr in parymids[i]:
                            # import pdb; pdb.set_trace()
                            pyr_pos_513 = pyr.get_rel_pos(7)[::-1]
                            cv2.circle(pred_instance, pyr_pos_513, 3, (0,0,255), -1)
                    parts_to_show.append(pred_instance)

            masked_imgs.append(np.hstack(parts_to_show))

        masked_show = np.vstack(masked_imgs)
        if tool=='cv2':
            cv2.imshow('Observation V5', masked_show)
            cv2.waitKey(10)
        elif tool=='plt':
            plt.figure('show', figsize=(16,12))
            plt.imshow(masked_show)
            plt.tight_layout()
            plt.show()

    def make_observation(self, images, sem_targets=None, ins_target=None,
            size_limit=None, s_image=True, s_sem_out=False, sematic_out=None,
            typical_size=513, s_sem_tar=False, s_ins_tar=False, s_ins_out=True,
            ins_out=None, ins_out_label=None, pyramids=None,
            s_run_mask=False, run_mask=None):
        show_dict = dict()
        N, _, img_h, img_w = images.shape
        out_img = []
        for bid in range(N):
            per_bid_show_img = []
            image_np = images[bid].permute(1,2,0).detach().cpu().numpy()
            image_np = ((image_np - image_np.min()) / (image_np.max() - image_np.min()))
            if s_image:
                per_bid_show_img.append(image_np)
            if s_sem_tar:
                sem_target_np = sem_targets[bid].detach().cpu().numpy()
                res_img = self.rend_on_image_v8(image_np, sem_target_np, range(19))
                per_bid_show_img.append(res_img)
            if s_sem_out:
                sematic_out_np = sematic_out[bid].max(0)[1].detach().cpu().numpy()
                res_img = self.rend_on_image_v8(image_np, sematic_out_np, range(19))
                per_bid_show_img.append(res_img)
            if s_ins_tar:
                ins_target_np = ins_target[bid].detach().cpu().numpy()
                ins_cats = [sem_targets[bid][ins_target[bid] == ins_idx].unique().item() for ins_idx in ins_target[bid].unique()]
                res_img = self.rend_on_image_v8(image_np, ins_target_np, ins_cats)
                per_bid_show_img.append(res_img)
            if s_ins_out:
                ins_out_np = ins_out[bid][0].max(0)[1].detach().cpu().numpy()
                res_img = self.rend_on_image_v8(image_np, ins_out_np, ins_out_label[bid])
                cv2.putText(res_img,'#:{}'.format(len(pyramids[bid])),(30,30),cv2.FONT_HERSHEY_COMPLEX,1,(1,0,0),1)
                for pyr in pyramids[bid]:
                    # import pdb; pdb.set_trace()
                    # mask_513 = pyr.get_mask(513)[0,0]
                    # accurate_pos_513 = tuple(torch.nonzero(mask_513 == mask_513.max())[0].tolist())
                    pyr_pos_513 = pyr.get_rel_pos(8)[::-1]
                    # cv2.circle(res_img, accurate_pos_513, 5, (1,0,0), -1)
                    cv2.circle(res_img, pyr_pos_513, 5, (1,0,0), -1)
                per_bid_show_img.append(res_img)

            # if s_run_mask:
            #     import pdb; pdb.set_trace()
            #     run_mask_np = run_mask

            for oimg in per_bid_show_img:
                if oimg.shape[:2] != (typical_size, typical_size):
                    import pdb; pdb.set_trace()
            if len(per_bid_show_img) < 8:
                out_img.append(np.hstack(per_bid_show_img))
            else:
                import pdb; pdb.set_trace()
        out = np.vstack(out_img)
        if size_limit:
            import pdb; pdb.set_trace()
        return out

        return show_dict

    def show_image(self, images, out_dict=None, show_via='cv2'):
        show_img = out_dict['show_img']
        if show_via == 'cv2':
            show_img = show_img[:,:,::-1]
            cv2.imshow('Observation CV2', show_img)
            cv2.waitKey(10)
        elif show_via == 'plt':
            plt.figure('show', figsize=(16,12))
            plt.imshow(show_img)
            plt.tight_layout()
            plt.show()
        elif show_via == 'output':
            return show_img[:,:,::-1]

    def forward(self, image, targets=None):
        output_dict = dict()
        N, _, img_h, img_w = image.shape
        sematic_info = self.sematic_bone(image)
        sematic_out = sematic_info[0]
        sematic_loss = self.sematic_criteron(sematic_out, targets['label'].long())
        sematic_info += (sematic_loss,)
        sematic_out, down_features, up_features, sematic_loss = sematic_info

        output_dict['loss_dict'] = {'sematic': sematic_loss.mean(), 'instance': 0}
        output_dict['sematic_out'] = sematic_out
        output_dict['targets'] = targets
        output_dict['ins_pyramids'] =[[] for _ in range(N)]
        output_dict['observe_images'] =[[] for _ in range(N)]

        if self.cfg.SOLVER.SEMATIC_ONLY:
            output_dict['show_img'] = self.make_observation(image, sematic_out=sematic_out,
                    sem_targets=targets['label'], ins_target=targets['instance'],
                    s_sem_tar=True, s_ins_tar=True, s_sem_out=True, )
            # self.show_image(image, output_dict, show_via='plt')
            return output_dict

        keep_down_features = self.keep_down(down_features[-1])
        down_features = down_features + keep_down_features
        lvl_sizes = [tuple(f.shape[-2:]) for f in down_features[::-1]]

        pyr_losses = []
        ins_out = []
        ins_pyramids = [[] for _ in range(N)]
        instance_categories = [11, 12, 13, 14, 15, 16, 17, 18]
        for bid in range(N):
            if self.training:
                ins_target_catch_pyr_id = [-1 for i in range(targets['instance'].unique().max()+1)]
                for sem_lvl in range(6):
                    lvl_size = lvl_sizes[sem_lvl]
                    inst_target = self.prepare_targets(targets, mode='instance', down_size=lvl_size)[bid]
                    sematic_target = self.prepare_targets(targets, mode='sematic', down_size=lvl_size)[bid]
                    sematic_lvl_resize = F.interpolate(sematic_out, lvl_size, mode='bilinear', align_corners=True)[bid]
                    for ins_tar_id in inst_target.unique():
                        target_cat = sematic_target[inst_target == ins_tar_id].unique().item()
                        if ins_target_catch_pyr_id[ins_tar_id] >= 0 or   \
                            target_cat not in instance_categories:
                            continue
                        else:
                            # if len(ins_pyramids[bid]) > 10:
                            #     continue
                            if (inst_target == ins_tar_id).sum() > 1:
                                pos_candis = torch.zeros_like(inst_target).float()
                                pos_candis[inst_target == ins_tar_id] = sematic_lvl_resize[target_cat][inst_target == ins_tar_id]
                                pos = tuple(torch.nonzero(pos_candis == pos_candis.max())[0].tolist())
                                # print('sem_lvl', sem_lvl, ins_tar_id, '\n', inst_target)
                                # import pdb; pdb.set_trace()
                            else:
                                pos = tuple(torch.nonzero(inst_target == ins_tar_id)[0].tolist())
                            new_pyr = InstancePyramid(pos, sem_lvl, lvl_sizes, ins_tar_id, target_cat)
                            ins_target_catch_pyr_id[ins_tar_id] = new_pyr.idx
                            ins_pyramids[bid].append(new_pyr)
                            # sematic_cat = sematic_lvl_resize[:, pos[0], pos[1]].argmax()
                            # if sematic_cat != target_cat:
                            #     print('Oops! it spots here to be {}, truth is {}.'.format(sematic_cat, target_cat))
                            #     import pdb; pdb.set_trace()
                    pyrs_this_lvl = [pyr for pyr in ins_pyramids[bid] if pyr.init_lvl == sem_lvl]
                    if pyrs_this_lvl:
                        self.compute_mask(sem_lvl, bid, down_features, pyrs_this_lvl)
                mask_stuff = sematic_out[[bid], :11]
                if len(ins_pyramids[bid]) == 0:
                    instance_loss = image.sum()* 0
                    mask_sem_ins = mask_stuff
                else:
                    # TODO: 合在一起来减少显存
                    # mask_ins_129 = torch.cat([pyr.get_mask_129() for pyr in ins_pyramids[bid]], dim=1)
                    # mask_ins = F.interpolate(mask_ins_129, 513, mode='bilinear', align_corners=True)
                    # mask_sem_ins = torch.cat((mask_stuff, mask_ins), dim=1)
                    mask_sem_ins_list = [mask_stuff]
                    for pyr in ins_pyramids[bid]:
                        mask_129 = pyr.get_mask_129()
                        mask_513 = F.interpolate(mask_129, 513, mode='bilinear', align_corners=True)
                        pyr.set_mask(513, mask_513)
                        mask_sem_ins_list.append(mask_513)
                    mask_sem_ins = torch.cat(mask_sem_ins_list, dim=1)

                    # target sem int
                    inst_target_idx = targets['label'][bid].clone()
                    inst_target_idx[inst_target_idx>=11] = -2
                    for i, pyr in enumerate(ins_pyramids[bid], start=11):
                        inst_target_idx[targets['instance'][bid]==pyr.ins_tar_id] = i
                    if -2 in inst_target_idx.unique():
                        # print('Warning! Uncovered instance mask left on its sematic mask!, sum={}'.format(
                        #     (inst_target_idx==-2).sum()))
                        # import pdb; pdb.set_trace()
                        inst_target_idx[inst_target_idx==-2] = -1
                    inst_target_idx = inst_target_idx[None]
                    instance_loss = self.sem_ins_criteron(mask_sem_ins, inst_target_idx)
                    # import pdb; pdb.set_trace()
            else:
                scratch_sematic_129 = F.interpolate(sematic_out, 129, mode='bilinear', align_corners=True)[bid]
                stuff_max_129 = scratch_sematic_129[:11].max(0)[0]
                for sem_lvl in range(6):
                    lvl_size = lvl_sizes[sem_lvl]

                    running_sem_ins_129 = torch.cat([scratch_sematic_129]+[pyr.get_mask_129()[0] for pyr in ins_pyramids[bid]], dim=0)
                    running_sem_ins_lvl = F.interpolate(running_sem_ins_129[None], lvl_size, mode='bilinear', align_corners=True)[0]
                    max_value, max_idx = running_sem_ins_lvl.max(0)
                    running_mask = (max_idx >= 11) * (max_idx < 19)
                    # 策略1： 新生只与11个sematic类作比较
                    pos_list = torch.nonzero(running_mask)
                    for pos_t in pos_list:
                        pos = tuple(pos_t.tolist())
                        # new_pos = tuple(torch.nonzero(max_value == max_value[running_mask].max())[0].tolist())
                        new_pyr = InstancePyramid(pos, sem_lvl, lvl_sizes, tar_cat=int(max_idx[pos]))
                        ins_pyramids[bid].append(new_pyr)
                        self.compute_mask(sem_lvl, bid, down_features, [new_pyr])

                        # import pdb; pdb.set_trace()
                        new_pyr_over_mask = new_pyr.get_mask_129()[0, 0] > stuff_max_129
                        if new_pyr_over_mask.sum() == 0:
                            # import pdb; pdb.set_trace()
                            continue
                        scratch_sematic_129 = scratch_sematic_129 * (1-new_pyr_over_mask.float())

                        # output_dict['show_img'] = self.make_observation(image, sematic_out=sematic_out,
                        #         sem_targets=targets['label'], ins_target=targets['instance'],
                        #         ins_out = ins_out, ins_out_label=None, pyramids=ins_pyramids,
                        #         s_sem_tar=False, s_ins_tar=True, s_sem_out=True, s_ins_out=False,
                        #         s_run_mask=True, run_mask=running_mask)
                        # self.show_image(image, output_dict, show_via='plt')

                mask_stuff = sematic_out[[bid], :11]
                if len(ins_pyramids[bid]) == 0:
                    instance_loss = image.sum()* 0
                    mask_sem_ins = mask_stuff
                else:
                    # TODO: 合在一起来减少显存
                    # mask_ins_129 = torch.cat([pyr.get_mask_129() for pyr in ins_pyramids[bid]], dim=1)
                    # mask_ins = F.interpolate(mask_ins_129, 513, mode='bilinear', align_corners=True)
                    # mask_sem_ins = torch.cat((mask_stuff, mask_ins), dim=1)
                    mask_sem_ins_list = [mask_stuff]
                    for pyr in ins_pyramids[bid]:
                        mask_129 = pyr.get_mask_129()
                        mask_513 = F.interpolate(mask_129, 513, mode='bilinear', align_corners=True)
                        pyr.set_mask(513, mask_513)
                        mask_sem_ins_list.append(mask_513)
                    mask_sem_ins = torch.cat(mask_sem_ins_list, dim=1)

                    # target sem int
                    inst_target_idx = targets['label'][bid].clone()
                    inst_target_idx[inst_target_idx>=11] = -2
                    for i, pyr in enumerate(ins_pyramids[bid], start=11):
                        inst_target_idx[targets['instance'][bid]==pyr.ins_tar_id] = i
                    if -2 in inst_target_idx.unique():
                        # print('Warning! Uncovered instance mask left on its sematic mask!, sum={}'.format(
                        #     (inst_target_idx==-2).sum()))
                        # import pdb; pdb.set_trace()
                        inst_target_idx[inst_target_idx==-2] = -1
                    inst_target_idx = inst_target_idx[None]
                    instance_loss = self.sem_ins_criteron(mask_sem_ins, inst_target_idx)
            pyr_losses.append(instance_loss)
            ins_out.append(mask_sem_ins)
        ins_out_label = [[i for i in range(11)]+[pyr.tar_cat for pyr in pyrs] for pyrs in ins_pyramids]
        output_dict['show_img'] = self.make_observation(image, sematic_out=sematic_out,
                sem_targets=targets['label'], ins_target=targets['instance'],
                ins_out = ins_out, ins_out_label=ins_out_label, pyramids=ins_pyramids,
                s_sem_tar=False, s_ins_tar=True, s_sem_out=True, s_ins_out=True)
        # self.show_image(image, output_dict, show_via='plt')

        output_dict['loss_dict'] = {'sematic': sematic_loss.mean(), 'instance': sum(pyr_losses)/N}
        output_dict['ins_pyramids'] =ins_pyramids

        return output_dict

class InstancePyramid():
    inst_count = 0
    def __init__(self, pos, init_lvl, lvl_sizes, ins_tar_id=-1, tar_cat=-1):
        self.idx = InstancePyramid.inst_count
        InstancePyramid.inst_count += 1
        self.pos = pos  # tuple required
        self.init_lvl = init_lvl
        self.lvl_sizes = lvl_sizes
        self.init_size = self.lvl_sizes[self.init_lvl]
        self.ins_tar_id = ins_tar_id
        self.tar_cat = tar_cat
        self.mask_129 = None
        self.masks = {}

        init_mask = torch.zeros(self.init_size)
        init_mask[self.pos] = 1
        cover_conv = nn.Conv2d(1,1,3,1,1, bias=False)
        cover_conv.weight.requires_grad = False
        cover_conv.weight.data.fill_(1)
        self.init_cover = cover_conv(init_mask[None,None])
        self.init_center_cover = cover_conv(init_mask[None,None])
        # center_mask = torch.ones(3,3)
        # if self.pos[0] <= 0:
        #     center_mask[:1, :] = 0
        # elif self.pos[0] >= self.init_size[0] - 1:
        #     center_mask[-1:, :] = 0
        # if self.pos[1] <= 0:
        #     center_mask[:, :1] = 0
        # elif self.pos[1] >= self.init_size[1] - 1:
        #     center_mask[:, -1:] = 0
        # self.center_mask = center_mask[None,None]

    def get_cover_mask_with_pyr_mask(self, lvl):
        f_mask = F.interpolate(self.init_cover, self.lvl_sizes[lvl], mode='nearest')[0,0]
        # pyr_mask = F.interpolate(self.center_mask, self.lvl_sizes[lvl-self.init_lvl], mode='nearest')
        pyr_mask = F.interpolate(self.center_mask, 2**(lvl-self.init_lvl+1)+1, mode='nearest')[0,0]
        return f_mask, pyr_mask

    def get_cover_mask(self, lvl):
        f_mask = F.interpolate(self.init_cover, self.lvl_sizes[lvl], mode='nearest')[0,0]
        # import pdb; pdb.set_trace()
        positive_pos = torch.nonzero(f_mask)
        pos_range = (positive_pos[:,0].min().item(), positive_pos[:,0].max().item()+1,
                    positive_pos[:,1].min().item(), positive_pos[:,1].max().item()+1)
        return f_mask, pos_range

    def set_mask(self, lvl, mask):
        self.masks[lvl] = mask

    def get_mask(self, lvl):
        return self.masks[lvl]

    def set_mask_129(self, mask):
        self.mask_129 = mask

    def get_mask_129(self):
        return self.mask_129

    def set_ins_target(self, lvl, ins_target):
        self.ins_target[lvl] = ins_target

    def get_ins_target(self, lvl):
        return self.ins_target[lvl]

    def set_mask_logits(self, lvl, mask_logits):
        self.mask_logits[lvl] = mask_logits

    def get_mask_logits(self, lvl):
        return self.mask_logits[lvl]

    def bind_target(self, idx):
        self.target_idx = idx

    def get_root_level_pos(self, pub_level):
        init_size = self.lvl_sizes[self.init_lvl]
        req_size = self.lvl_sizes[pub_level]

        h = (self.pos[0].float() / init_size[0] * req_size[0]).round().long()
        w = (self.pos[1].float() / init_size[1] * req_size[1]).round().long()

        return (h.item(), w.item())

    def get_rel_pos(self, pub_level):
        init_size = self.lvl_sizes[self.init_lvl]
        req_size = self.lvl_sizes[pub_level]

        h = round((self.pos[0]+0.5) / init_size[0] * req_size[0]-0.5)
        w = round((self.pos[1]+0.5) / init_size[1] * req_size[1]-0.5)

        return (h, w)

    def shared_gaussian_masks(self):
        # feature_scales = [7, 13, 25, 50, 100, 200]
        xs = torch.arange(7*4)
        ys = torch.arange(7*4).view(-1,1)
        ln2 = torch.Tensor(2.0, requires_grad=False).log()
        # shared_gaussian_mask = (-4*ln2*((xs.float()-7*2+1)**2+(ys.float()-7*2+1)**2)/7**2).exp()
        shared_gaussian_mask = (-ln2*((xs.float()-7*2+1)**2+(ys.float()-7*2+1)**2)/7**2).exp()
        return shared_gaussian_mask

    def init_gaussian_masks(self):
        xs = torch.arange(7*4)
        ys = torch.arange(7*4).view(-1,1)
        # ln2 = torch.tensor(2.0, requires_grad=False).log()
        ln2 = torch.Tensor([2.0]).log()[0]
        gaussian_mask_28 = (-ln2*((xs.float()-7*2+1)**2+(ys.float()-7*2+1)**2)/7**2).exp()

        init_gaussian_mask = torch.zeros(self.level_sizes[self.init_level])

        src_x0 = max(13-self.pos[0], 0)
        src_y0 = max(13-self.pos[1], 0)
        src_x1 = min(13+self.level_sizes[self.init_level][0]-self.pos[0], 28)
        src_y1 = min(13+self.level_sizes[self.init_level][1]-self.pos[1], 28)
        res_x0 = max(0, self.pos[0]-13)    #+1?
        res_y0 = max(0, self.pos[1]-13)
        res_x1 = res_x0+src_x1-src_x0
        res_y1 = res_y0+src_y1-src_y0

        init_gaussian_mask[res_x0:res_x1, res_y0:res_y1] = gaussian_mask_28\
            [src_x0:src_x1, src_y0:src_y1]

        return init_gaussian_mask

    def get_indicate_mask(self, feature):
        indicate = torch.zeros_like(feature[0,0])
        indicate[tuple(self.pos)] = 1.0
        mask = torch.ones_like(feature[0,0]) * feature[0,0][tuple(self.pos)]
        return torch.stack((indicate, mask))

class Trainer(object):
    def __init__(self, cfg, output_dir):
        self.cfg = cfg
        self.output_dir = output_dir
        self.logger = logging.getLogger("MaskPyramid")
        self.tbSummary = TensorboardSummary(output_dir)
        self.writer = self.tbSummary.create_summary()
        self.model = MaskPyramids(cfg)
        self.device = torch.device(cfg.MODEL.DEVICE)
        self.model.to(self.device)
        self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(cfg)

        train_params = [{'params': self.model.keep_down.parameters(), 'lr': cfg.SOLVER.BASE_LR},
            {'params': self.model.cat_conv.parameters(), 'lr': cfg.SOLVER.BASE_LR},]
        self.optimizer = torch.optim.SGD(train_params, momentum=cfg.SOLVER.MOMENTUM,
                weight_decay=cfg.SOLVER.WEIGHT_DECAY, nesterov=cfg.SOLVER.NESTEROV)

        self.evaluator = Evaluator(self.nclass)
        self.scheduler = LR_Scheduler(cfg.SOLVER.SCHEDULE_TYPE, cfg.SOLVER.BASE_LR,
            cfg.SOLVER.EPOCHES, len(self.train_loader))
        self.start_epoch = 0
        self.best_pred = 0.0
        self.meters = {'start_time': time.time(),
            'total_iters': cfg.SOLVER.EPOCHES*len(self.train_loader)}

        # log_gpu_stat(self.logger)

    def load_weights(self, path=None, subdict='model', continue_train=False):
        state_dict = torch.load(path if path else self.cfg.MODEL.WEIGHT)
        if subdict:
            weights = state_dict[subdict]
            self.model.load_state_dict(weights)
            if self.optimizer is not None and 'optimizer' in state_dict.keys():
                self.optimizer.load_state_dict(state_dict["optimizer"])
            if 'best_pred' in state_dict.keys():
                self.best_pred = state_dict["best_pred"]
            if continue_train and 'epoch' in state_dict.keys():
                self.start_epoch = state_dict["epoch"]
        else:
            self.model.load_state_dict(state_dict)

    def mask_inst_img(self, images_np, sematic_a_image_np, output, idx, level=-1, label_mode=1, insts_a_image=None):
        pyramids = output['all_pyramids'][idx]
        inst_target = None
        if level>=0:
            import pdb; pdb.set_trace()
        bg_and_masks = [F.interpolate(pyramids[0].get_mask(3)[:,[0]], 513, mode='bilinear', align_corners=True)]
        pred_labels = [0]
        sematic_out = output['sematic_out'][idx]
        for pyr in pyramids:
            # import pdb; pdb.set_trace()
            pos = tuple(pyr.pos.tolist())
            largest_mask = pyr.get_mask(3)
            mask_513 = F.interpolate(largest_mask[:,[1]], 513, mode='bilinear', align_corners=True)
            # mask_513_map = mask_513[0].max(0)[1]
            bg_and_masks.append(mask_513)
            # prepare label
            scaled_sematic_out = F.interpolate(sematic_out[None], pyr.init_size, mode='bilinear', align_corners=True)
            label = scaled_sematic_out[0].max(0)[1][pos[0], pos[1]].item()
            pred_labels.append(label)

        # import pdb; pdb.set_trace()
        bg_and_masks_np = torch.cat(bg_and_masks).squeeze(1).max(0)[1].detach().cpu().numpy()
        inst_output = self.rend_on_image(images_np, bg_and_masks_np, pred_labels)

        return  inst_output, inst_target

    def rend_on_image_v8(self, image_np, masks_np, labels_list):
        # label 从0开始，0 表示unlabeled
        alpha = 0.5
        color_bias = 1
        # import pdb; pdb.set_trace()
        # print('masks_np unique',np.unique(masks_np), 'labels_list', labels_list)
        class_count = [0 for _ in range(20)]
        colors = np.array([
            # [0, 100, 0],
            [128, 64, 128],
            [244, 35, 232],
            [70, 70, 70],
            [102, 102, 156],
            [190, 153, 153],
            [153, 153, 153],
            [250, 170, 30],
            [220, 220, 0],
            [107, 142, 35],
            [152, 251, 152],
            [0, 130, 180],
            [220, 20, 60],
            [255, 0, 0],
            [0, 0, 142],
            [0, 0, 70],
            [0, 60, 100],
            [0, 80, 100],
            [0, 0, 230],
            [119, 11, 32]])
        # import pdb; pdb.set_trace()
        masked = image_np.copy()
        # import pdb; pdb.set_trace()
        for i, label in enumerate(labels_list):
            if label == -1:
                continue
            mask_idx = np.nonzero(masks_np == i)
            masked[mask_idx[0], mask_idx[1], :] *= 1.0 - alpha
            color = (colors[label]+class_count[label]*color_bias)/255.0
            # try:
            #     color = (colors[label]+class_count[label]*color_bias)/255.0
            # except:
            #     import pdb; pdb.set_trace()

            masked[mask_idx[0], mask_idx[1], :] += alpha * color
            # import pdb; pdb.set_trace()
            contours, hierarchy = cv2.findContours(
                (masks_np == i).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
            )
            # masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), -1)
            masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), 1)
            class_count[label] += 1

        return masked

    def show_image_v8(self, images, output, targets):
        N,C,H,W = images.shape
        masked_imgs = []
        origin_image_ON = True
        ori_sematic_target_ON = False
        ori_instance_target_ON = True
        car_only_ON = False
        pred_sematic_ON = True
        pred_instance_ON = True
        pred_stat_ON = True
        for i, image in enumerate(images):
            parts_to_show = []
            images_np = image.permute(1,2,0).detach().cpu().numpy()
            images_np = ((images_np - images_np.min()) / (images_np.max() - images_np.min()))[...,::-1]
            if origin_image_ON:
                parts_to_show.append(images_np)
            semat_a_target, insts_a_target = targets['label'][i], targets['instance'][i]

            # targets prepared in forwards
            # out_one = output['in_dict_watches'][i]
            # target_fwd_label = (out_one['target_levels']['labels']+1).tolist()
            sematic_target = output['targets']['label'][i]
            instacne_target = output['targets']['instance'][i]
            # sematic_out_513 = output['sematic_out'][6]
            sematic_out_513 = output['sematic_out']

            if ori_sematic_target_ON:
                sematic_target_np = sematic_target.detach().cpu().numpy()
                ori_sematic_target = self.rend_on_image_v8(images_np, sematic_target_np, range(19))
                parts_to_show.append(ori_sematic_target)

            if ori_instance_target_ON:
                instacne_target_np = instacne_target.detach().cpu().numpy()
                ins_cats = [sematic_target[instacne_target == ins_idx].unique().item() for ins_idx in instacne_target.unique()]
                if car_only_ON:
                    ins_cats = [13 if label==13 else -1 for label in ins_cats]
                ori_instacne_target = self.rend_on_image_v8(images_np, instacne_target_np, ins_cats)
                parts_to_show.append(ori_instacne_target)

            # import pdb; pdb.set_trace()
            # prediction    - sematic
            if pred_sematic_ON:
                sematic_predict_np = sematic_out_513[i].max(0)[1].detach().cpu().numpy()
                sematic_predict = self.rend_on_image_v8(images_np, sematic_predict_np, range(19))
                parts_to_show.append(sematic_predict)

            if not self.cfg.SOLVER.SEMATIC_ONLY:
                parymids = output['ins_pyramids']
                # prediction    - instance
                if pred_instance_ON:
                    # import pdb; pdb.set_trace()
                    instance_predict = torch.cat([pyr.get_mask(513) for pyr in  parymids[i]], dim=1) if len(parymids[i]) else torch.empty((1,1,513,513))
                    instance_predict_513 = F.interpolate(instance_predict, 513, mode='bilinear', align_corners=True)
                    instance_predict_np = instance_predict_513[0].max(0)[1].detach().cpu().numpy()
                    ins_cats_predict = [pyr.tar_cat for pyr in  parymids[i]]
                    # import pdb; pdb.set_trace()
                    pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, ins_cats_predict)
                    # pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, range(19))
                    if pred_stat_ON:
                        cv2.putText(pred_instance,'#:{}'.format(len(parymids[i])),(30,30),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1)
                        for pyr in parymids[i]:
                            # import pdb; pdb.set_trace()
                            pyr_pos_513 = pyr.get_rel_pos(7)[::-1]
                            cv2.circle(pred_instance, pyr_pos_513, 3, (0,0,255), -1)
                parts_to_show.append(pred_instance)

            masked_imgs.append(np.hstack(parts_to_show))

        masked_show = np.vstack(masked_imgs)
        cv2.imshow('Observation V5', masked_show)
        cv2.waitKey(10)

    def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        end = time.time()
        for i, sample in enumerate(self.train_loader):
            image, label, instance = sample
            image, label, instance = image.to(self.device), label.to(self.device), instance.to(self.device)
            target = {'label': label, 'instance': instance}
            self.scheduler(self.optimizer, i, epoch, self.best_pred)
            self.optimizer.zero_grad()
            output_dict = self.model(image, target)
            loss_dict = output_dict['loss_dict']
            losses = sum(loss for loss in loss_dict.values())
            losses.backward()
            self.optimizer.step()
            batch_time = time.time() - end
            end = time.time()
            train_loss += losses.item()

            if i % 20 == 0 or i == len(self.train_loader) -1:
                curr_iter = epoch*len(self.train_loader)+ i + 1
                sepent_time = time.time() - self.meters['start_time']
                eta_seconds = sepent_time * self.meters['total_iters'] / curr_iter - sepent_time
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if self.cfg.SOLVER.SEMATIC_ONLY:
                    self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|SematicLoss:{:2.4}|Class:{:2.4}').format(
                        epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
                        eta_string, losses.item(), loss_dict['sematic'].item()))
                else:
                    self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|SematicLoss:{:2.4}|Class:{:2.4}').format(
                        epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
                        eta_string, losses.item(), loss_dict['sematic'].item()))

            if self.cfg.SOLVER.SHOW_IMAGE and i % 10 == 0:
                self.model.show_image(image, output_dict)
                if i // 10 % 10 == 0:
                    observe_out = self.model.show_image(image, output_dict, show_via='output')
                    cv2.imwrite(self.output_dir+"/ObserveOut_E_{}_I{}.png".format(epoch, i), observe_out*255)
        self.writer.add_scalar('train/loss_epoch', train_loss, epoch)

    def validation(self, epoch):
        self.model.eval()
        self.evaluator.reset()
        test_loss = 0.0
        tbar = tqdm(self.val_loader, desc='\r')
        for i, sample in enumerate(tbar):
            image, label, instance = sample
            image, label, instance = image.to(self.device), label.to(self.device), instance.to(self.device)
            target = {'label': label, 'instance': instance}
            with torch.no_grad():
                output_dict = self.model(image, target)
            # import pdb; pdb.set_trace()
            sematic_out = output_dict['sematic_out']
            pred = sematic_out.data.cpu().numpy()
            pred = np.argmax(pred, axis=1)
            # Add batch sample into evaluator
            self.evaluator.add_batch(label.detach().cpu().numpy(), pred)

            if self.cfg.SOLVER.SHOW_IMAGE and i % 5 == 0:
                self.model.show_image(image, output_dict)

        # Fast test during the training
        # import pdb; pdb.set_trace()
        Acc = self.evaluator.Pixel_Accuracy()
        Acc_class = self.evaluator.Pixel_Accuracy_Class()
        mIoU = self.evaluator.Mean_Intersection_over_Union()
        FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
        self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
        self.writer.add_scalar('val/mIoU', mIoU, epoch)
        self.writer.add_scalar('val/Acc', Acc, epoch)
        self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
        self.writer.add_scalar('val/fwIoU', FWIoU, epoch)

        self.logger.info('Evalueat report: mIoU: {:3.4}| Acc: {:3.4}| Acc_class: {:3.4}| fwIoU: {:3.4}| previousBest: {:3.4}|'.format(
            mIoU, Acc, Acc_class, FWIoU, float(self.best_pred)
        ))
        if mIoU > self.best_pred:
            is_best = True
            self.best_pred = mIoU
            save_data = {}
            save_data["epoch"] = epoch + 1
            save_data["best_pred"] = self.best_pred
            save_data["model"] = self.model.state_dict()
            if self.optimizer is not None:
                save_data["optimizer"] = self.optimizer.state_dict()
            if self.scheduler is not None:
                save_data["scheduler"] = self.scheduler.__dict__
            torch.save(save_data, self.output_dir+"/model_Epoch_{}.pth".format(epoch))

def log_gpu_stat(logger):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
    used_mem = (meminfo.used / 1024) /1024
    pynvml.nvmlShutdown()
    gpu_info = subprocess.check_output(["nvidia-smi"])
    logger.info(("\nThe pid of current job is {} and {}, the used memory before we run is {}MB,"+\
        " the <nvidia-smi> shows:\n{}").format(os.getpid(),os.getppid(), used_mem, gpu_info.decode("utf-8")))

def main():
    parser = argparse.ArgumentParser(description="PyTorch MaskPyramid Training for Cityscapes")
    parser.add_argument(        # args.config
        "--config",
        default="maskpyramid/config/gallery/MaskPyramidV5.yml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(        # args.opts
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    args = parser.parse_args()

    cfg.merge_from_file(args.config)
    # cfg.merge_from_list(['SOLVER.SHOW_IMAGE', False])
    cfg.merge_from_list(['DATALOADER.BATCH_SIZE_TRAIN', 2])
    # cfg.merge_from_list(['DATALOADER.NUM_WORKERS', 0])
    # cfg.merge_from_list(['SOLVER.SEMATIC_ONLY', True])
    cfg.merge_from_list(['MULTI_RUN', False])
    cfg.merge_from_list(['MODEL.WEIGHT', 'run/mpc_v11.1_1/Every_5_model_Epoch_40.pth'])
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    extra_exp_name = 'sematic' if cfg.SOLVER.SEMATIC_ONLY else ''
    experiment_name = os.path.basename(__file__).split('.py')[0]
    experiment_name += '_'+extra_exp_name if extra_exp_name else ''
    if cfg.MULTI_RUN:
        run_names = glob.glob(cfg.OUTPUT_DIR + '/*')
        matched_nums = []
        for name in run_names:
            matched_nums += re.findall(r'%s_(\d+)'%experiment_name, name)
        experiment_name += '_' + str(max([int(n) for n in matched_nums])+1) if matched_nums else '_1'
    output_dir = os.path.join(cfg.OUTPUT_DIR, experiment_name)

    mkdir(output_dir)
    with open(os.path.join(output_dir, experiment_name+'_config.yml'), 'w') as f:
        f.write(cfg.dump())
    logger = setup_logger("MaskPyramid", output_dir, 0, experiment_name+"_log.txt")  # get_rank())
    logger.info(args)
    logger.info("Collecting env info (might take some time)")
    logger.info("\n" + collect_env_info())
    logger.info("Loaded and archive configuration file at '{}'".format(os.path.join(output_dir, experiment_name+'_config.yml')))
    log_gpu_stat(logger)

    trainer = Trainer(cfg, output_dir)
    logger.info('Starting Epoch:{} / Total Epoches:{}'.format(trainer.start_epoch, cfg.SOLVER.EPOCHES))
    if cfg.MODEL.WEIGHT:
        trainer.load_weights(cfg.MODEL.WEIGHT)

    for epoch in range(trainer.start_epoch, cfg.SOLVER.EPOCHES):
        # trainer.training(epoch)
        trainer.validation(epoch)
        log_gpu_stat(logger)
        if epoch == 5 or epoch % 20 == 0:
            save_data = {}
            save_data["epoch"] = epoch + 1
            save_data["best_pred"] = -1
            save_data["model"] = trainer.model.state_dict()
            if trainer.optimizer is not None:
                save_data["optimizer"] = trainer.optimizer.state_dict()
            if trainer.scheduler is not None:
                save_data["scheduler"] = trainer.scheduler.__dict__
            torch.save(save_data, trainer.output_dir+"/Every_5_model_Epoch_{}.pth".format(epoch))

    trainer.writer.close()

if __name__ == "__main__":
    main()

# v11.2 完成了validation，但是由于train val 不统一导致的大片覆盖问题下一版要解决
# v11.1 train 编码完成； 无validation; observe 完善
# v11 backbone 与 maskpyramid 分开， 直接利用好的backbone weight
# v10_5 instance 引入
# v10_5_resunet_v3 没有aspp
# v10_5_resunet_v2_sematic_1 层次加深到 feature map size 最小为4
# v10.2: imshow 展示pyr num 与实际的instance target
# v10: Unet 下，不再多层输出，仅使用最终mask 通过放缩来匹配筛选各级mask

# v9: resnet 换成 Unet

# v8: 新的策略：
# 首先使用sematic segmentation 的 cross entropy 得出19个类的logits，
# 然后，对于可数的类，找未instance化的极值，初始化为instance，同时，instance 在计算时侵蚀（×0）掉该class logis
# 注意， 这个侵蚀， training时用target 侵蚀， testing 用instance 高于其它类别的值侵蚀
# 要有个while循环来把所有sematic 部分instance化


#Программа для сборки пекеджей баз валют
from builder_modules import *
from tkinter import *

#Создаем приложение
tk = Tk()
tk.title("Currency package builder")
tk.resizable(0, 0)
tk.iconbitmap("title_image.ico")
canvas = Canvas(tk, width=300, height=50,  highlightthickness=0)
canvas.pack()

#Создаем поле ввода пути к папкам скриптов для выбранного устройства путем введения его названия
path_title = Label(tk, text='Device model:', fg='green', font=('Verdana', 10)).pack()
path_entry = Entry(tk, width=10, fg='blue', font=('Verdana', 10)).pack()

#Указываем путь к файлам сборки необходимого устройства
path = r'C:\SVN_Projects\NOTEBASE\D820\VARIANT'

#Необходимое расширение
extension = '.lua'

#Создаем поле ввода валюты
currency_title = Label(tk, text='Currency name:', fg='green', font=('Verdana', 10)).pack()
currency_entry = Entry(tk, width=5, fg='blue', font=('Verdana', 10)).pack()

#Вибираем валюту с измененной базой
#*************************
currency = '{ \"UZS\"'
#*************************

#Находим файлы с расширением .lua
list_of_lua_files = list_of_fils(path, extension)

#Находим список скриптов для запуска
scripts_for_launching = files_for_start(list_of_lua_files, currency)

#Добавляем кнопку для сборки пекеджей
builder_button = Button(tk, text='build', command=lambda scrip=scripts_for_launching: scripts_running(scrip))
builder_button.pack()

tk.mainloop()


import sys
import json

# takes the selected user dictionary entry as parameter from the calling funtion/file (or main?)
user_info=json.loads(sys.argv[1])# !!!DOES NOT WORK!!! Alternate to be checked

amount=input("Enter amount to dispense:")
amount=int(amount)
for acc_balance in user_info:
    if (amount <= user_info[acc_balance]):
        print("Please collect cash, Thanks for using our ATM!")
    else:
        print("Insufficient balance")
user_info[acc_balance]-=amount #how to return this value to the calling funtion to update acc_balance in the original dictionary?


"""Event Study Abnormal Returns

- key developments CAR, BHAR and post-event drift

Terence Lim
License: MIT
"""
import numpy as np
import pandas as pd
from pandas import DataFrame, Series
import matplotlib.pyplot as plt
import os
import time
from finds.database import SQL
from finds.busday import BusDay
from finds.structured import PSTAT, CRSP, Benchmarks
from finds.backtesting import EventStudy
from settings import settings
LAST_DATE = settings['crsp_date']

ECHO = True
sql = SQL(**settings['sql'], echo=ECHO)
user = SQL(**settings['user'], echo=ECHO)
bd = BusDay(sql)
keydev = PSTAT(sql, bd)
crsp = CRSP(sql, bd, rdb=None)
bench = Benchmarks(sql, bd)
eventstudy = EventStudy(user, bench, LAST_DATE)
outdir = os.path.join(settings['images'], 'events')

# event window parameters
end = 20201201
beg = 19890101  # 20020101
minobs = 250
left, right, post = -1, 1, 21

# str formatter to pretty print event and role description given their id's
eventformat = lambda e, r: "{event} ({eventid}) {role} [{roleid}]".format(
    event=keydev.event_[e], eventid=e, role=keydev.role_[r], roleid=r)
events = sorted(keydev.event_.keys())   # list of eventid's
roles = sorted(keydev.role_.keys())     # list of roleid's

## Helpers to merge events and crsp, and screen stock universe

# to lookup prevailing exchange and share codes by permno and date
shrcd = crsp.build_lookup('permno', 'shrcd')
exchcd = crsp.build_lookup('permno', 'exchcd')

def event_pipeline(eventstudy, mincap=300000, **arg):
    """helper to merge keydev events and crsp, and screen stock universe"""
    df = keydev.get_linked(
        dataset='keydev',
        date_field='announcedate',
        fields=['keydevid',
                'keydeveventtypeid',
                'keydevtoobjectroletypeid'],
        where=(f"announcedate >= {arg['beg']} and announcedate <= {arg['end']}"
               f" and keydeveventtypeid = {arg['eventid']} "
               f" and keydevtoobjectroletypeid = {arg['roleid']}"))\
               .drop_duplicates(['permno','announcedate'])\
               .set_index(['permno','announcedate'], drop=False)

    stk = arg['stocks'].get_many(
        dataset='daily',
        permnos=df['permno'],
        date_field='date',
        dates=arg['stocks'].bd.offset(df['announcedate'], left-1),
        fields=['prc', 'shrout']).fillna(0)

    df['cap'] = (stk['prc'].abs() * stk['shrout']).values
    df['exchcd'] = [exchcd(r.permno, r.date) for r in stk.itertuples()]
    df['shrcd'] = [shrcd(r.permno, r.date) for r in stk.itertuples()]
    r = (df['cap'].gt(mincap) &              # require cap > $300M
         df['exchcd'].isin([1,2,3]) &        # primary exchange
         df['shrcd'].isin([10,11])).values   # domestic common stocks
    rows = eventstudy(crsp, df[r], left, right, post, 'announcedate')
    return df.loc[rows.to_records(index=False).tolist()]  # restrict df to rows

## Compute abnormal returns of all events
# %%capture
tic = time.time()
for i, eventid in enumerate(events):
    for roleid in roles:
        # retrieve all observations of this eventid, roleid
        df = event_pipeline(eventstudy, stocks=crsp, beg=beg, end=end,
                            eventid=eventid, roleid=roleid,
                            left=left, right=right, post=post)
        if len(df) < minobs:  # require min number of events
            continue

        # retrieve announcement window returns
        r = eventstudy(crsp, df, left, right, post, 'announcedate')
        if r['date'].nunique() < minobs:  # require min number of dates
            continue

        # compute both BHAR and CAR averages, plot and save
        bhar = eventstudy.fit(car=False, name='bhar')
        car = eventstudy.fit(car=True, name='car')
        eventstudy.write(label=f"{eventid}_{roleid}")
        s = pd.concat([bhar, car], axis=1).T
        print(eventformat(eventid, roleid))
        print(s.to_string(float_format='%.4f', index=False))
        print()
        fig, axes = plt.subplots(2, 1, clear=True, num=1, figsize=(10,12))
        eventstudy.plot(title=eventformat(eventid, roleid),
                        vline=right, ax=axes[0], name='bhar')
        eventstudy.plot(title='', vline=right, ax=axes[1], name='car')
        if outdir:
            plt.savefig(os.path.join(outdir, f"{eventid}_{roleid}.jpg"))
    print('Elapsed:', time.time()-tic, 'secs')

## Summarize BHAR's of all events, by 3-day event window abnormal returns

df = eventstudy.read(name='bhar')\
               .set_index('permno').sort_values('window', ascending=False)
dx = DataFrame(df.index.str.split('_').to_list()).astype(int)
df.index = pd.MultiIndex.from_frame(dx).set_names(['eventid','roleid'])

df['event'] = keydev.event_[df.index.get_level_values(0)].values
df['role'] = keydev.role_[df.index.get_level_values(1)].values

mindays = (df['days']>1000).values
print(df[mindays].iloc[:10].drop(columns='name')\
      .to_string(formatters={'effective':'{:.0f}'.format}, float_format='%.4f',
                 index=False))
print(df[mindays].iloc[::-1].iloc[:10].drop(columns='name')\
      .to_string(formatters={'effective':'{:.0f}'.format}, float_format='%.4f',
                 index=False))
print(df[mindays].iloc[:10].drop(columns=['name'])\
      .to_latex(index=False, formatters={'effective':'{:.0f}'.format}))
print(df[mindays].iloc[::-1].iloc[:10].drop(columns=['name'])\
      .to_latex(index=False, formatters={'effective':'{:.0f}'.format}))
print(df.sort_values('post_t').to_string(float_format='%.4f'))
df.drop(columns='name')

## Show single plots for each of three events
eventid, roleid = 80, 1
eventid, roleid = 26, 1
df = event_pipeline(eventstudy, stocks=crsp, eventid=eventid, roleid=roleid,
                    beg=beg, end=end, left=left, right=right, post=post)
bhar = eventstudy.fit(car=False)
fig, ax = plt.subplots(clear=True, num=1, figsize=(10,6))
eventstudy.plot(title=eventformat(eventid, roleid), vline=right, ax=ax)

## show single plot by market cap and half-period
midcap = 20000000
for i, (eventid, roleid) in enumerate([[50,1], [83,1]]):
    #eventid, roleid = 50, 1
    #eventid, roleid = 83, 1
    df = event_pipeline(eventstudy, stocks=crsp, eventid=eventid, roleid=roleid,
                        beg=beg, end=end, left=left, right=right, post=post)
    halfperiod = np.median(df['announcedate'])
    sample = {'[FirstHalf]': df['announcedate'].ge(halfperiod).values,
              '[SecondHalf]': df['announcedate'].lt(halfperiod).values,
              '[Large]': df['cap'].ge(midcap).values,
              '[Small]': df['cap'].lt(midcap).values,
              '': None}
    for ifig, (label, rows) in enumerate(sample.items()):
        fig, ax = plt.subplots(clear=True, num=1+ifig, figsize=(5,6))
        bhar = eventstudy.fit(rows=rows, car=False)
        eventstudy.plot(title=eventformat(eventid, roleid) + ' ' + label,
                        drift=True, ax=ax, c=f"C{i*5+ifig}")
        plt.savefig(os.path.join(outdir, label + f"{eventid}_{roleid}.jpg"))

for i, (eventid, roleid) in enumerate([[80,1], [26,1]]):
    #eventid, roleid = 50, 1
    #eventid, roleid = 83, 1
    df = event_pipeline(eventstudy, stocks=crsp, eventid=eventid, roleid=roleid,
                        beg=beg, end=end, left=left, right=right, post=post)
    halfperiod = np.median(df['announcedate'])
    sample = {'[FirstHalf]': df['announcedate'].ge(halfperiod).values,
              '[SecondHalf]': df['announcedate'].lt(halfperiod).values,
              '[Large]': df['cap'].ge(midcap).values,
              '[Small]': df['cap'].lt(midcap).values,
              '': None}
    for ifig, (label, rows) in enumerate(sample.items()):
        fig, ax = plt.subplots(clear=True, num=1+ifig, figsize=(5,6))
        bhar = eventstudy.fit(rows=rows, car=False)
        eventstudy.plot(title=eventformat(eventid, roleid) + ' ' + label,
                        drift=False, ax=ax, c=f"C{i*5+ifig}")
        plt.savefig(os.path.join(outdir, label + f"{eventid}_{roleid}.jpg"))

#plt.show()

## Show Max Order Statistic and Bonferroni Adjustment
import statsmodels.api as sm
import scipy
from pandas.api import types
class MaxStat:
    """Max Order Statistic probability distributions"""
    def __init__(self, dist=scipy.stats.norm, n=None, **params):
        self.dist_ = dist
        self.params_ = params
        self.n = n

    def cdf(self, z, n=None):
        """cdf for max order statistic"""
        return [self.cdf(y, n) for y in z] if types.is_list_like(z)\
            else self.dist_.cdf(z, **self.params_)**(n or self.n)

    def pdf(self, z, n=None):
        """cdf for max order statistic"""
        n = n or self.n
        return [self.pdf(y, n) for y in z] if types.is_list_like(z)\
            else self.dist_.pdf(z, **self.params_) * n * self.cdf(z, n=n-1)

    def ppf(self, z, n=None):
        """inverse cdf for max order statistic"""
        return [self.ppf(y, n) for y in z] if types.is_list_like(z)\
            else self.dist_.ppf(z, **self.params_)**(n or self.n)

    def pvalue(self, z, n=None):
        """z-value for max order statistic"""
        return [self.pvalue(y, n) for y in z] if types.is_list_like(z)\
            else 1 - (self.dist_.cdf(z, **self.params_)**(n or self.n))

    def zvalue(self, p, n=None):
        """z-value for max order statistic"""
        return [self.zvalue(y, n) for y in z] if types.is_list_like(p)\
            else self.dist_.ppf((1-p)**(1/(n or self.n)), **self.params_)

    def bonferroni(self, p, n=None):
        """corrected z-value with with Bonferroni adjustment"""
        return [self.bonferroni(y, n) for y in z] if types.is_list_like(p)\
            else self.dist_.ppf((1-(p/n)), **self.params_)

# Display order statistics of folded normal
y = eventstudy.read(name='bhar')['post_t'].values
Z = max(y)
n = len(y)
print(f"Events tested={n}, Max z-value={Z:.4f}:")
maxstat = MaxStat(scipy.stats.foldnorm, c=0)

p = Series({n: maxstat.pvalue(Z, n=n) for n in
            sorted([n] + [1, 2, 30, 60, 120, 250, 1000, 1600])}, name='pvalue')
print(f"\nMax order statistic p-value(z={Z:.2f}) by sample size:")
print(DataFrame(p).T.to_string(float_format='%.4f'))

P=0.05
print(f"\nRejection Region(p-value={P:.2f}) by sample size:")
zb = Series({n: maxstat.bonferroni(P, n=n)
             for n in [1, 2, 20, 100, 1000, 100000, 1000000]},name='max-order')
zc = Series({n: maxstat.zvalue(P, n=n)
             for n in [1, 2, 20, 100, 1000, 100000, 1000000]},name='bonferroni')
print(pd.concat([zc, zb], axis=1).round(3).to_latex())

# Plot CDF of Max Order Statistic
X = np.linspace(0, 6, 600)
df = DataFrame(index=X)
df = pd.concat([Series(data=[maxstat.cdf(x, n) for x in X], index=X, name=n)
                for n in [1, 2, 20, 100, 1000, 100000, 1000000]], axis=1)
fig, ax = plt.subplots(clear=True, num=1, figsize=(10,6))
df.plot(ax=ax, title='CDF of Max-Order Statistic by Sample Size',
        xlabel='z-value', ylabel='Cumulative Distribution Function')
ax.axvline(1.96, c='grey', alpha=0.5)
ax.annotate("1.96", xy=(1.96, 1))
plt.savefig(os.path.join(outdir, 'cdf.jpg'))
plt.show()

## Show distribution of post-event t-values
import scipy
import seaborn as sns
y = eventstudy.read(name='bhar')['post_t']
fig, axes = plt.subplots(1, 2, num=1, clear=True, figsize=(10,6))

ax = sns.distplot(y, kde=False, hist=True, ax=axes[0], color='C0')
bx = ax.twinx()
x = np.linspace(*plt.xlim(), 100)
bx.plot(x, scipy.stats.norm.pdf(x), color="C1")
ax.set_title('Post-event drift and Normal Density', fontsize=10)
ax.set_xlabel('Post-event drift t-value')

ax = sns.distplot(abs(y), kde=False, hist=True, ax=axes[1], color='C1')
bx = ax.twinx()
x = np.linspace(*plt.xlim(), 100)
bx.plot(x, scipy.stats.foldnorm.pdf(x, 0), color="C1")
ax.set_title('Abs(post-event drift) and Folded Normal Density', fontsize=10)
ax.set_xlabel('Abs(post-event drift) t-value')
plt.savefig(os.path.join(outdir, 'hist.jpg'))

DataFrame(y.describe()).T