import argparse
import os

import torch

from net.models import LeNet
from net.quantization import apply_weight_sharing
import util

parser = argparse.ArgumentParser(description='This program quantizes weight by using weight sharing')
parser.add_argument('model', type=str, help='path to saved pruned model')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--output', default='saves/model_after_weight_sharing.ptmodel', type=str,
                    help='path to model output')
args = parser.parse_args()

use_cuda = not args.no_cuda and torch.cuda.is_available()

# Define the model
model = torch.load(args.model)
print('accuracy before weight sharing')
util.test(model, use_cuda)

# Weight sharing
apply_weight_sharing(model)
print('accuacy after weight sharing')
util.test(model, use_cuda)

# Save the new model
os.makedirs('saves', exist_ok=True)
torch.save(model, args.output)


import re

import wx

from wx import GridSizer
from wx.lib.agw.supertooltip import SuperToolTip

from boaui.units import area, charge, inertia, length, mass, pressure, volume, tnt, density, torque

from .label import SmartLabel
from . import LayoutDimensions, SmartToolTip

from ..units import KEY_IMPERIAL, KEY_METRIC
from ..units.acceleration import AccelerationUnit
from ..units.angle import AngleUnit
from ..units.area_density import AreaDensityUnit
from ..units.area import AreaUnit
from ..units.charge import ChargeUnit
from ..units.density import DensityUnit
from ..units.force import ForceUnit
from ..units.inertia import InertiaUnit
from ..units.length import LengthUnit
from ..units.linear_density import LinearDensityUnit
from ..units.linear_pressure import LinearPressureUnit
from ..units.mass import MassUnit
from ..units.pressure import PressureUnit
from ..units.time import TimeUnit
from ..units.tnt import TntUnit
from ..units.torque import TorqueUnit
from ..units.velocity import VelocityUnit
from ..units.volume import VolumeUnit

__author__ = 'Joeny'

class SmartTextBox(wx.TextCtrl):
    """
    Create a smarter text box that could capture keys and process them
    to see if the format is correct.

    The validation method goes through three process:

        1. OnChar(): Capture ony the key character that are necessary.
        2. wx.EVT_TEXT: Validate that the input is actually a number.
        3. Validate(): Check against the tolerance level.

    """
    def __init__(self, parent, key_up=None, message=None, enabled_message='',
                 disabled_messages=None, disabled_index=None, value=None, enable=None,
                 helptext=None, required=False,
                 normal=(255, 255, 255), format_error=(228, 115, 115), range_error=(244, 67, 54),
                 data=None, *args, **kwargs):
        """
        Constructor

        :param parent: parent ui
        :param key_up: bind key up handler
        :param message: add in tooltip message
        :param enabled_message: message once the box is enabled
        :param disabled_messages: list of array messages
        :param disabled_index: index of the which messages to display
        :param value: initial value for smart box
        :param enable: enable box
        :param helptext: add in context help button
        :param required: requirement
        :param normal: rgb
        :param format_error: rgb
        :param range_error: rgb
        :param data: used to hold any unique data that can be assessed later
        :param args:
        :param kwargs:
        """
        wx.TextCtrl.__init__(self, parent, *args, **kwargs)

        if value is not None:
            self.Value = str(value)

        self.keys = kwargs.get('keys', {})
        self.parent = parent
        self.data = data

        if key_up:
            self.Bind(wx.EVT_KEY_UP, key_up, self)

        self.tooltip = None
        if message:
            self.tooltip = wx.ToolTip(message)
            self.SetToolTip(self.tooltip)

        self.enabled_message = enabled_message
        self.disabled_messages = disabled_messages

        if disabled_index is None and self.disabled_messages:
            self.disabled_index = 0
        else:
            self.disabled_index = disabled_index

        if helptext:
            self.SetHelpText(helptext)

        self.required = required

        self.color_normal = normal
        self.color_format_error = format_error
        self.color_range_error = range_error

        if enable is not None:
            self.Enable(enable)

    def clear(self):
        self.Clear()

    @property
    def min(self):
        """
        Return the minimum value.

        :return: minimum value
        """
        return self.keys.get('min')

    @min.setter
    def min(self, value):
        """
        Set the minimum value.

        :param value:
        """
        self.keys['min'] = value

    @property
    def max(self):
        """
        Return the maximum value.

        :return: return max value
        """
        return self.keys.get('max')

    @max.setter
    def max(self, value):
        """
        Set the maximum value.

        :param: value
        """
        self.keys['max'] = value

    def set_value(self, value, fmt=None):
        # type: (object, object) -> object
        """
        Set the textbox value

        :param value: text
        :return:
        """
        if value is not None:
            if fmt:
                self.Value = fmt%(value)
            else:
                self.Value = str(value)
        else:
            self.Value = ""

    def get_value(self, key=None):
        """
        Get the value

        :param key:
        :return:
        """
        val = self.GetValue()

        if key is not None:
            # When key is strike we capture.
            digit = chr(key)

            pos = self.GetInsertionPoint()
            if pos == len(val):
                val += digit
            else:
                val = val[:pos] + digit + val[pos:]

        return val

    def Enable(self, *args, **kwargs):
        """
        On enable, clean data if needed.

        :param args:
        :param kwargs:
        """
        wx.TextCtrl.Enable(self, *args, **kwargs)

        if self.disabled_messages:
            if self.Value in self.disabled_messages:
                self.Value = self.enabled_message

    def Disable(self, *args, **kwargs):
        """
        On disable, add message if needed.

        :param args:
        :param kwargs:
        """
        wx.TextCtrl.Disable(self, *args, **kwargs)

        if self.disabled_messages:
            self.set_disable_message()

    def set_normal_color(self):
        """
        Set normal color.

        """
        self.SetBackgroundColour(self.color_normal)
        self.Refresh()

    def set_format_error_color(self):
        """
        Set format error color.

        """
        self.SetBackgroundColour(self.color_format_error)
        self.Refresh()

    def set_range_error_color(self):
        """
        Set range error color.

        """
        self.SetBackgroundColour(self.color_range_error)
        self.Refresh()

    def set_disable_message(self):
        """
        Set disable message.

        :return:
        """
        self.Value = self.disabled_messages[self.disabled_index]

    def check_requirement(self):
        """
        Check if the textbox has value

        :return:
        """
        if self.required:
            if self.Enabled:
                if self.get_value() is None:

                    # Set error box.
                    if hasattr(self, 'set_range_error_color'):
                        self.set_format_error_color()

                    return False

        if hasattr(self, 'set_normal_color'):
            self.set_normal_color()

        # If not required, than return true.
        return True

class SmartComboBox(wx.ComboBox):
    """
    Smart ComboBox is used for units conversion.

    """
    def __init__(self, parent, data=None, style=wx.CB_READONLY, value='', message=None, unit=None, unit_system=None,
                 enabled_message='', disabled_messages=None, disabled_index=None, enable=None,
                 helptext=None, required=False, *args, **kwargs):
        """
        Constructor

        :param parent: parent panel or frame
        :param data: list of values
        :param style: combobox style
        :param value: display value
        :param message: tooltip message
        :param unit: Unit object
        :param unit_system: 'imperial' or 'metric'
        :param enabled_message: enable message
        :param disabled_messages: disable message
        :param disabled_index:
        :param enable: enable combobox
        :param helptext: add in context help
        :param required:
        :param args:
        :param kwargs:
        :return:
        """
        wx.ComboBox.__init__(self, parent, style=style, *args, **kwargs)

        self.convert = None
        self.unit_system = unit_system
        self.unit = unit

        if data:
            self.AppendItems(data)

        if value:
            self.Value = value

        self.tooltip = None
        if message:
            self.tooltip = wx.ToolTip(message)
            self.SetToolTip(self.tooltip)

        self.previous_index = 0
        self.enabled_message = enabled_message
        self.disabled_messages = disabled_messages

        if disabled_index is None and self.disabled_messages:
            self.disabled_index = 0
        else:
            self.disabled_index = disabled_index

        if unit:
            # If unit is passed in, activate it.
            self.activate()

        self.current_dropbox_selection = None

        self.Bind(wx.EVT_COMBOBOX_DROPDOWN, self.on_dropdown_open, self)

        if helptext:
            self.SetHelpText(helptext)

        self.required = required

        if enable is not None:
            self.Enable(enable)

    def bind_dropdown(self, handle):
        """
        Bind dropdown event to handle.

        """
        self.Bind(wx.EVT_COMBOBOX, handle)

    def Enable(self, *args, **kwargs):
        """
        On enable, clean data if needed.

        :param args:
        :param kwargs:
        """
        wx.ComboBox.Enable(self, *args, **kwargs)

        if self.disabled_messages:
            if self.Value in self.disabled_messages:
                for index, label in enumerate(self.Strings):
                    if label in self.disabled_messages:
                        self.Delete(index)

                self.SetSelection(self.previous_index)

    def Disable(self, *args, **kwargs):
        """
        On disable, add message if needed.

        :param args:
        :param kwargs:
        """
        wx.ComboBox.Disable(self, *args, **kwargs)

        if self.disabled_messages:
            self.previous_index = self.GetCurrentSelection()
            self.Append(self.disabled_messages[self.disabled_index])
            self.SetSelection(self.GetCount() - 1)

    def on_dropdown_open(self, event=None):
        """
        Event handler to store the current selection

        :param:
        """
        self.current_dropbox_selection = self.GetCurrentSelection()

    def is_selection_change(self):
        """
        Check if the dropbox selection different from the previous selection before the dropbox is open.

        :return: boolean
        """
        if self.current_dropbox_selection is self.GetSelection():
            return False
        else:
            return True

    def append(self, label, obj):
        """
        Append data into combobox.

        :param label: title
        :param obj: object data
        :return:
        """
        self.Append(label, obj)

    def set_selection_by_data(self, value):
        """
        Set the selection given the data input.

        :param value:
        :return:
        """
        for index, text in enumerate(self.Strings):
            if self.HasClientData():
                if self.GetClientData(index) == value:
                    self.SetSelection(index)

                    # Leave loop
                    return

    def get_data(self):
        """
        Get the data.

        :return:
        """
        if self.GetSelection() == -1:
            return None
        else:
            return self.GetClientData(self.GetSelection())

    def set_value(self, value):
        # type: (object) -> object
        """
        Set the value

        :param value: string
        :return:
        """
        self.Value = str(value)

    def get_value(self):
        """
        Get the combobox value

        :return:
        """
        return self.Value

    def activate(self):
        """
        Activate Units.

        :return:
        """
        self.Clear()
        self.AppendItems(self.unit.get_list())
        self.SetSelection(self.unit.get_default_selection())
        self.convert = self.unit.get_conversion_factor

    def activate_acceleration(self, *args, **kwargs):
        """
        Activate acceleration unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = AccelerationUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_angle(self, *args, **kwargs):
        """
        Activate angle unit

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = AngleUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_area_density(self, *args, **kwargs):
        """
        Activate area density unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = AreaDensityUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_area(self, *args, **kwargs):
        """
        Activate area unit.

        :param kwargs:
        """
        self.unit = AreaUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_charge(self, *args, **kwargs):
        """
        Activate charge weight.

        :param kwargs:
        """
        self.unit = ChargeUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_density(self, *args, **kwargs):
        """
        Activate density unit.

        :param args:
        :param kwargs:
        """
        self.unit = DensityUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_force(self, *args, **kwargs):
        """
        Active force unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = ForceUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_inertia(self, *args, **kwargs):
        """
        Activate Inertia unit.

        :param args:
        :param kwargs:
        """
        self.unit = InertiaUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_length(self, *args, **kwargs):
        """
        Activate length unit.

        :param args:
        :param kwargs:
        """
        self.unit = LengthUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_linear_density(self, *args, **kwargs):
        """
        Activate linear density unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = LinearDensityUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_linear_pressure(self, *args, **kwargs):
        """
        Activate linear pressure unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = LinearPressureUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_mass(self, *args, **kwargs):
        """
        Activate mass units.

        :param kwargs:
        """
        self.unit = MassUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_pressure(self, *args, **kwargs):
        """
        Activate pressure unit.

        :param kwargs:
        """
        self.unit = PressureUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_time(self, *args, **kwargs):
        """
        Activate time unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = TimeUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_tnt(self, *args, **kwargs):
        """
        Activate tnt unit.

        :param args:
        :param kwargs:
        """
        self.unit = TntUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_torque(self, *args, **kwargs):
        """
        Activate Torque unit.

        :param args:
        :param kwargs:
        """
        self.unit = TorqueUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_velocity(self, *args, **kwargs):
        """
        Activate Velocity unit.

        :param args:
        :param kwargs:
        :return:
        """
        self.unit = VelocityUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def activate_volume(self, *args, **kwargs):
        """
        Activate volume unit.

        :param args:
        :param kwargs:
        """
        self.unit = VolumeUnit(*args, **kwargs)

        if kwargs.get('unit_list'):
            unit_list = kwargs.get('unit_list')

            self.unit.metric_list = unit_list['metric']
            self.unit.imperial_list = unit_list['imperial']

        self.unit.unit_system = self.unit_system
        self.activate()

    def get_factor(self, origin, destination):
        """
        Get the factor.

        :param origin: origin unit
        :param destination: destination unit
        """
        return self.convert(origin, destination)

    def check_requirement(self):
        """
        Check if the textbox has value

        :return:
        """
        if self.required:
            if self.Enabled:
                if self.get_value() is None:
                    return False
                else:
                    return True
            else:
                # If textbox is not active, than it's not required.
                return True
        else:
            # If not required, than return true.
            return True

class SmartInputLayout(wx.BoxSizer):
    """
    Create the horizontal layout of smart textbox.

    /---------------- OVERALL WIDTH ----------------------/
    |                                                     |
    |                                                     |

    *******************************************************     ----/
    *    *     * *                     * *              * *         |
    *******************************************************         |
    *    *     * *                     * *              * *
    *    * (1) * *         (2)         * *     (3)      * *     OVERALL HEIGHT
    *    *     * *                     * *              * *
    *******************************************************         |
    *    *     * *                     * *              * *         |
    *******************************************************     ----/

    """

    MAKE_VERTICAL_STRETCHABLE = 1

    def __init__(self, parent, max=None, min=None, layout=None, label=None, *args, **kwargs):
        """
        Constructor.

        :param parent:
        :param width:
        :param max: maximum value for the textbox
        :param min: minimum value for the textbox
        :param layout:
        :param label: pass in wx.Label or SmartLabel
        :param args:
        :param kwargs:
        """
        wx.BoxSizer.__init__(self, wx.VERTICAL)

        self.components = []
        self.parent = parent
        self.hsizer = None
        self._next_id = 0
        self.INDEX_LABEL = None
        self.INDEX_TEXTBOX = None
        self.INDEX_POSTBOX = None
        self.INDEX_COMBOBOX = None
        self.INDEX_BUTTON = None

        if layout:
            self.layout = layout
        else:
            self.layout = LayoutDimensions()

            self.layout.calculate()

        # Add in the label.
        if label:
            self.label = label
        elif kwargs.get('name'):
            self.label = SmartLabel(self.parent, label=kwargs.get('name'))
        else:
            self.label = None

        self.min = min
        self.max = max

        # Set minimum size.
        size = self.GetSize()
        size.Height = self.layout.overall_height
        self.SetMinSize(size)

    def rename(self, name=None):
        """
        Rename

        :param name:
        :return:
        """
        self.label = wx.StaticText(self.parent, label=name)

    @property
    def next_id(self):
        """

        :return:
        """
        nid = self._next_id
        self._next_id += 1
        return nid

    @property
    def label(self):
        """

        :return:
        """
        if self.INDEX_LABEL is None:
            return None

        return self.components[self.INDEX_LABEL]

    @label.setter
    def label(self, value):
        """

        :param value:
        :return:
        """
        if value is None:
            return

        self.INDEX_LABEL = self.next_id
        self.components.append(value)

    @property
    def textbox(self):
        """

        :return:
        """
        if self.INDEX_TEXTBOX is None:
            return None

        return self.components[self.INDEX_TEXTBOX]

    @textbox.setter
    def textbox(self, value):
        """

        :param value:
        :return:
        """
        if value is None:
            return

        self.INDEX_TEXTBOX = self.next_id
        self.components.append(value)

    @property
    def postbox(self):
        """

        :return:
        """
        if self.INDEX_POSTBOX is None:
            return None

        return self.components[self.INDEX_POSTBOX]

    @postbox.setter
    def postbox(self, value):
        """

        :param value:
        :return:
        """
        if value is None:
            return

        self.INDEX_POSTBOX = self.next_id
        self.components.append(value)

    @property
    def combobox(self):
        """

        :return:
        """
        if self.INDEX_COMBOBOX is None:
            return None

        return self.components[self.INDEX_COMBOBOX]

    @combobox.setter
    def combobox(self, value):
        """

        :param value:
        :return:
        """
        if value is None:
            return

        self.INDEX_COMBOBOX = self.next_id
        self.components.append(value)

    def do_layout(self):
        """
        Do Layout.

        :return:
        """

        # Start with the vertical margin.
        self.AddSpacer(self.layout.top)

        # Move from left to right.
        self.hsizer = wx.BoxSizer(wx.HORIZONTAL)
        # self.hsizer.SetMinSize(wx.Size(self.layout.overall_width, self.layout.height))

        self.hsizer.AddSpacer(self.layout.left)

        for id in range(0, len(self.components)):
            """
            wx.BoxSizer.Add(window, proportion=0, flag=0, border=0, userData=None)
                Append a child to the sizer

                :param window: a window, a spacer or another sizer to be added to the sizer.  Its initial size
                    (either set explicitly by the user or calculated internally) is interpreted as the minimal and
                    in many cases also the initial size.
                :param proportion: (int) the parameter is used in wx.BoxSizer to indicate if a child of a sizer can
                    change its size in the main orientation of the wx.BoxSizer - where 0 stands for non changeable
                    and a value of more than zero is interpreted relative to the value of other children of the
                    same wx.BosSizer.  For example, you might have a horizontal wx.BoxSizer with three children,
                    two of which are supposed to change their size with the sizer.  Then the two stretchable
                    windows would get a value of 1 each to make item grow and shrink equally with the sizer's
                    horizontal dimension.
                :param flag: (int): combinations of flags affecting sizer's behavior
                :param border: (int): determines the border width, if the flag parameter is set to include any
                    border flag
                :param userData: (object) allows an extra object to be attached to the sizer item, for use in
                    derived classes when sizing information
            """
            self.components[id].SetMinSize(self.layout.get_size(id))

            self.hsizer.AddSpacer(self.layout.interior)
            self.hsizer.Add(self.components[id],
                            self.layout.stretch_factor[id],
                            wx.ALL | wx.EXPAND,
                            self.layout.border_width[id])

        # Add blank space if no component exists.
        for id_blank in range(id+1, len(self.layout.widths)):
            self.hsizer.AddSpacer(self.layout.interior)

            blank_label = wx.StaticText(self.parent, label="")
            blank_label.SetMinSize(self.layout.get_size(id_blank))

            self.hsizer.Add(blank_label,
                            self.layout.stretch_factor[id_blank],
                            wx.ALL | wx.EXPAND,
                            self.layout.border_width[id_blank])

        self.hsizer.AddSpacer(self.layout.right)

        self.Add(self.hsizer, 1, wx.EXPAND | wx.ALL, 0)
        self.AddSpacer(self.layout.bottom)

    def add(self, item, proportion=0, flag=0, border=0, userData=None):
        """
        Appends a child item to the sizer.

        :param item: The item can be one of three kind of objects:

            * window: A wx.Window to be managed by the sizer. Its minimal size (either set explicitly by the user or
                calculated internally when constructed with wx.DefaultSize) is interpreted as the minimal size to use
                when laying out item in the sizer. This is particularly useful in connection with
                wx.Window.SetSizeHints.

            * sizer: The (child-)sizer to be added to the sizer. This allows placing a child sizer in a sizer and thus
                to create hierarchies of sizers (for example a vertical box as the top sizer and several horizontal
                boxes on the level beneath).

            * size: A wx.Size or a 2-element sequence of integers that represents the width and height of a spacer to
                be added to the sizer. Adding spacers to sizers gives more flexibility in the design of dialogs;
                imagine for example a horizontal box with two buttons at the bottom of a dialog: you might want to
                insert a space between the two buttons and make that space stretchable using the proportion value and
                the result will be that the left button will be aligned with the left side of the dialog and the right
                button with the right side - the space in between will shrink and grow with the dialog.

        :param proportion: Although the meaning of this parameter is undefined in wx.Sizer, it is used in wx.BoxSizer
            to indicate if a child of a sizer can change its size in the main orientation of the wx.BoxSizer - where 0
            stands for not changeable and a value of more than zero is interpreted relative (a proportion of the total)
            to the value of other children of the same wx.BoxSizer. For example, you might have a horizontal
            wx.BoxSizer with three children, two of which are supposed to change their size with the sizer. Then the
            two stretchable windows should each be given proportion value of 1 to make them grow and shrink equally
            with the sizer's horizontal dimension. But if one of them had a proportion value of 2 then it would get a
            double share of the space available after the fixed size items are positioned.

            (type int)

        :param flag: This parameter can be used to set a number of flags which can be combined using the binary OR
            operator |. Two main behaviours are defined using these flags. One is the border around a window: the
            border parameter determines the border width whereas the flags given here determine which side(s) of the
            item that the border will be added. The other flags determine how the sizer item behaves when the space
            allotted to the sizer changes, and is somewhat dependent on the specific kind of sizer used.

            * wx.TOP
            * wx.BOTTOM
            * wx.LEFT
            * wx.RIGHT
            * wx.ALL

            * wx.EXPAND

            * wx.SHAPED

            * wx.FIXED_MINSIZE

            * wx.ALIGN_CENTER
            * wx.ALIGN_LEFT
            * wx.ALIGN_RIGHT
            * wx.ALIGN_TOP
            * wx.ALIGN_BOTTOM
            * wx.ALIGN_CENTER_VERTICAL
            * wx.ALIGN_CENTER_HORIZONTAL

            (type int)

        :param border: Determines the border width, if the flag parameter is set to include any border flag.

            (type int)

        :param userData: Allows an extra object to be attached to the sizer item, for use in derived classes when
            sizing information is more complex than the proportion and flag will allow for.

            (type=PyObject)
        """
        self.Add(item, proportion, flag, border, userData)

    def add_stretch_spacer(self, prop=1):
        """
        Add a stretchable spacer.

        :param prop:
        :return:
        """
        self.AddStretchSpacer(prop=prop)

    def add_spacer(self, size):
        """
        Add a spacer that is (size, size) pixels.

        :param size:
        :return:
        """
        self.AddSpacer(size)

    def fit(self, window):
        """
        Tell the sizer to resize the window to match the sizer's minimal size. This is commonly done in the constructor
        of the window itself in order to set its initial size to match the needs of the children as determined by the
        sizer. Returns the new size.

        For a top level window this is the total window size, not the client size.

        :param window:
        :return:
        """
        self.Fit(window)

    def enable(self):
        """
        Must inherit enable input layout.

        """
        pass

    def disable(self):
        """
        Must inherit disable input layout.

        """
        pass

    def validate(self):
        """
        Must inherit validate().

        """
        pass

    def check_requirement(self):
        """
        Check requirement

        :return:
        """
        requirement_satisfy = True

        for item in self.components:
            if hasattr(item, 'check_requirement'):
                if item.check_requirement() is False:
                    requirement_satisfy = False

        return requirement_satisfy

class SmartButton(wx.Button):
    """
    Smarter Button Class

    """
    def __init__(self, parent, label='', evt_button=None, message=None, enable=None, helptext=None, *args, **kwargs):
        """
        Constructor.

        :param parent:
        :param label:
        :param evt_button:
        :param message:
        :param enable:
        :param helptext:
        :param args:
        :param kwargs:
        :return:
        """
        wx.Button.__init__(self, parent, label=label, *args, **kwargs)

        if evt_button:
            self.Bind(wx.EVT_BUTTON, evt_button)

        self.tooltip = None
        if message:
            self.tooltip = wx.ToolTip(message)
            self.SetToolTip(self.tooltip)

        if helptext:
            self.SetHelpText(helptext)

        if enable is not None:
            self.Enable(enable)

class SmartCheckBox(wx.CheckBox):
    """
    **Smarter CheckBox**

    """
    def __init__(self, parent, id=-1, label='', evt_click=None, message=None, enable=None, helptext=None,
                 *args, **kwargs):
        """
        Constructor.

        :param parent:
        :param id:
        :param label:
        :param evt_click:
        :param message:
        :param helptext:
        :param args:
        :param kwargs:
        :return:
        """
        wx.CheckBox.__init__(self, parent, id=id, label=label, *args, **kwargs)

        self.tooltip = None
        if message:
            self.tooltip = wx.ToolTip(message)
            self.SetToolTip(self.tooltip)

        if evt_click:
            self.Bind(wx.EVT_CHECKBOX, evt_click)

        if helptext:
            self.SetHelpText(helptext)

        if enable is not None:
            self.Enable(enable)

    def bind_click(self, handle):
        """
        Bind check box click.
        :param handle:
        :return:
        """
        self.Bind(wx.EVT_CHECKBOX, handle)

    def get_value(self):
        """
        Return the true/false

        :return:
        """
        return self.Value

    def set_value(self, value):
        # type: (object) -> object
        self.SetValue(value)


# 位1的个数
# https://leetcode-cn.com/leetbook/read/top-interview-questions-easy/xn1m0i/

class Solution:
    def hammingWeight(self, n: int) -> int:
        count = 0
        while n is not 0:
            count += n & 1
            n = n >> 1
        return count

r = Solution().hammingWeight(0b1111)
print(r)


from cavefinder.support.cstruct import *

# ---------------------
# | Dos Header        |
# ---------------------
# | Pe Signature      |
# ---------------------
# | COFF Header       |
# *********************
# | Optional Header   |
# *********************
# | Section Table     |
# ---------------------
# | Mappable sections |
# ---------------------

#
MZ_MAGIC = 0x4D5A
MZ_CIGAM = 0x5A4D

PE_MAGIC = 0x50450000
PE_CIGAM = 0x00004550

class DosHeader(object, metaclass=CStruct):
    e_magic = USHORT  # 00: MZ Header signature
    e_cblp = USHORT  # 02: Bytes on last page of file
    e_cp = USHORT  # 04: Pages in file
    e_crlc = USHORT  # 06: Relocations
    e_cparhdr = USHORT  # 08: Size of header in paragraphs
    e_minalloc = USHORT  # 0a: Minimum extra paragraphs needed
    e_maxalloc = USHORT  # 0c: Maximum extra paragraphs needed
    e_ss = USHORT  # 0e: Initial (relative) SS value
    e_sp = USHORT  # 10: Initial SP value
    e_csum = USHORT  # 12: Checksum
    e_ip = USHORT  # 14: Initial IP value
    e_cs = USHORT  # 16: Initial (relative) CS value
    e_lfarlc = USHORT  # 18: File address of relocation table
    e_ovno = USHORT  # 1a: Overlay number
    e_res = "8s"  # 1c: Reserved words
    e_oemid = USHORT  # 24: OEM identifier (for e_oeminfo)
    e_oeminfo = USHORT  # 26: OEM information=0 e_oemid specific
    e_res2 = "20s"  # 28: Reserved words
    e_lfanew = UINT  # 3c: Offset to extended header

    def __init__(self, stream: io.RawIOBase):
        self.e_magic = unpack_type(USHORT, stream.read(2))
        stream.seek(stream.tell() - sizeof(USHORT))

        if self.e_magic == MZ_CIGAM or MZ_MAGIC:
            self.unpack_from_io(stream)
        else:
            raise TypeError("Not a valid PE (Invalid DosHeader)")

    def __str__(self):
        return '\n'.join(["Dos Header:",
                          "Magic:                       {e_magic:#x}",
                          "Byte on last page:           {e_cblp:#x}",
                          "Page in file:                {e_cp:#x}",
                          "Relocations:                 {e_crlc:#x}",
                          "Size of header:              {e_cparhdr:#x}",
                          "Min alloc:                   {e_minalloc:#x}",
                          "Max alloc:                   {e_maxalloc:#x}",
                          "ss:                          {e_ss:#x}",
                          "sp:                          {e_sp:#x}",
                          "Checksum:                    {e_csum:#x}",
                          "ip:                          {e_ip:#x}",
                          "cs:                          {e_cs:#x}",
                          "File address reloc table:    {e_lfarlc:#x}",
                          "Overlay number:              {e_ovno:#x}",
                          "OEM id:                      {e_oemid:#x}",
                          "OEM info:                    {e_oeminfo:#x}",
                          "Offset to extended header:   {e_lfanew:#x}"]).format(**self.__dict__)

    @property
    def endianness(self):
        return 'big' if self.e_magic == MZ_MAGIC else 'little'

COFF_MACHINE_UNKNOWN = 0  # unknown
COFF_MACHINE_I386 = 0x014c  # Intel 386.
COFF_MACHINE_R3000 = 0x0162  # MIPS little-endian, 0x160 big-endian
COFF_MACHINE_R4000 = 0x0166  # MIPS little-endian
COFF_MACHINE_R10000 = 0x0168  # MIPS little-endian
COFF_MACHINE_WCEMIPSV2 = 0x0169  # MIPS little-endian WCE v2
COFF_MACHINE_ALPHA = 0x0184  # Alpha_AXP
COFF_MACHINE_POWERPC = 0x01F0  # IBM PowerPC Little-Endian
COFF_MACHINE_SH3 = 0x01a2  # SH3 little-endian
COFF_MACHINE_SH3E = 0x01a4  # SH3E little-endian
COFF_MACHINE_SH4 = 0x01a6  # SH4 little-endian
COFF_MACHINE_ARM = 0x01c0  # ARM Little-Endian
COFF_MACHINE_THUMB = 0x01c2
COFF_MACHINE_AMD64 = 0x8664
COFF_MACHINE_IA64 = 0x0200  # Intel 64
COFF_MACHINE_MIPS16 = 0x0266  # MIPS
COFF_MACHINE_MIPSFPU = 0x0366  # MIPS
COFF_MACHINE_MIPSFPU16 = 0x0466  # MIPS
COFF_MACHINE_ALPHA64 = 0x0284  # ALPHA64

class COFFHeader(object, metaclass=CStruct):
    machine = USHORT
    nsections = USHORT
    timestamp = UINT
    ptr_to_symtable = UINT
    nsym = UINT
    size_opheader = USHORT
    characteristics = USHORT

    def __init__(self, stream: io.RawIOBase, endianness):
        self.unpack_from_io(stream, endianness)

    def __str__(self):
        return "\n".join(["Machine:                 {machine:#x} (%s)",
                          "Number of sections:      {nsections}",
                          "Timestamp:               {timestamp}",
                          "Ptr to symbol table:     {ptr_to_symtable:#x}",
                          "Number of symbol:        {nsym}",
                          "Size of optional header: {size_opheader}",
                          "Characteristics:         {characteristics:#x}"]).format(**self.__dict__) % self.machine_str

    @property
    def machine_str(self):
        val = {
            COFF_MACHINE_UNKNOWN: "unknown",
            COFF_MACHINE_I386: "Intel 386",
            COFF_MACHINE_R3000: "MIPS little-endian",
            COFF_MACHINE_R4000: "MIPS little-endian",
            COFF_MACHINE_R10000: "MIPS little-endian",
            COFF_MACHINE_WCEMIPSV2: "MIPS little-endian WCE v2",
            COFF_MACHINE_ALPHA: "Alpha_AXP",
            COFF_MACHINE_POWERPC: "PowerPC little-endian",
            COFF_MACHINE_SH3: "SH3 little-endian",
            COFF_MACHINE_SH3E: "SH3E little-endian",
            COFF_MACHINE_SH4: "SH4 little-endian",
            COFF_MACHINE_ARM: "ARM little-endian",
            COFF_MACHINE_THUMB: "THUMB",
            COFF_MACHINE_AMD64: "AMD 64",
            COFF_MACHINE_IA64: "Intel 64",
            COFF_MACHINE_MIPS16: "MIPS",
            COFF_MACHINE_MIPSFPU: "MIPS",
            COFF_MACHINE_MIPSFPU16: "MIPS",
            COFF_MACHINE_ALPHA64: "ALPHA64"
        }

        return "Unknown: %02x" % self.machine if self.machine not in val else val[self.machine]

IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x10b
IMAGE_NT_OPTIONAL_HDR64_MAGIC = 0x20b
IMAGE_ROM_OPTIONAL_HDR_MAGIC = 0x107

IMAGE_SUBSYSTEM_UNKNOWN = 0x0
IMAGE_SUBSYSTEM_NATIVE = 0x1
IMAGE_SUBSYSTEM_WINDOWS_GUI = 0x2
IMAGE_SUBSYSTEM_WINDOWS_CUI = 0x3
IMAGE_SUBSYSTEM_OS2_CUI = 0x5
IMAGE_SUBSYSTEM_POSIX_CUI = 0x7
IMAGE_SUBSYSTEM_NATIVE_WINDOWS = 0x8
IMAGE_SUBSYSTEM_WINDOWS_CE_GUI = 0x9

class OptionalHeader(object, metaclass=CStruct):
    magic = USHORT
    major_linker_version = UCHAR
    minor_linker_version = UCHAR
    size_of_code = UINT
    size_initialized_data = UINT
    size_uninitialized_data = UINT
    address_entry_point = UINT
    base_code = UINT
    image_base = ULONGLONG
    section_alignment = UINT
    file_alignment = UINT
    major_osversion = USHORT
    minor_osversion = USHORT
    major_image_version = USHORT
    minor_image_version = USHORT
    major_subsystem_version = USHORT
    minor_subsystem_version = USHORT
    win32_version = UINT
    size_image = UINT
    size_headers = UINT
    checksum = UINT
    subsystem = USHORT
    dll_characteristics = USHORT
    size_stack_reserve = ULONGLONG
    size_stack_commit = ULONGLONG
    size_heap_reserve = ULONGLONG
    size_heap_commit = ULONGLONG
    loader_flags = UINT
    number_rva_and_sizes = UINT

    def __init__(self, stream: io.RawIOBase, endianness):
        self.unpack_from_io(stream, endianness)

    def subsytem_str(self):
        val = {
            IMAGE_SUBSYSTEM_UNKNOWN: "Unknown",
            IMAGE_SUBSYSTEM_NATIVE: "Native",
            IMAGE_SUBSYSTEM_WINDOWS_GUI: "Windows GUI",
            IMAGE_SUBSYSTEM_WINDOWS_CUI: "Console",
            IMAGE_SUBSYSTEM_OS2_CUI: "OS/2 Console",
            IMAGE_SUBSYSTEM_POSIX_CUI: "Posix Console",
            IMAGE_SUBSYSTEM_NATIVE_WINDOWS: "diver 9x native",
            IMAGE_SUBSYSTEM_WINDOWS_CE_GUI: "Windows CE",
        }
        return "Unknown: %02x" % self.subsystem if self.subsystem not in val else val[self.subsystem]

    @property
    def wordsz(self):
        if self.magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC:
            return 32
        elif self.magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC:
            return 64
        return 0

class PeSectionHeader(object, metaclass=CStruct):
    name = "8s"
    physaddr_or_virtsize = UINT
    virtual_addr = UINT
    size_rawdata = UINT
    ptr_rawdata = UINT
    ptr_relocations = UINT
    ptr_linenumbers = UINT
    nrelocations = USHORT
    nlinenumbers = USHORT
    characteristics = UINT

    def __init__(self, stream: io.RawIOBase, endianness):
        self.unpack_from_io(stream, endianness)

    def __str__(self):
        return "\n".join(["Section Header",
                          "Name:                    {name}",
                          "Virtual address:         {virtual_addr:#x}",
                          "Size of raw data:        {size_rawdata:#x}",
                          "Pointer to raw data:     {ptr_rawdata:#x}",
                          "Pointer to relocations:  {ptr_relocations:#x}",
                          "Pointer to line numbers: {ptr_linenumbers:#x}",
                          "Number of relocations:   {nrelocations:#x}",
                          "Number of line numbers:  {nlinenumbers:#x}",
                          "Characteristics:         {characteristics:#x}"]).format(**self.__dict__)

class PEHeader(object):
    def __init__(self, stream: io.RawIOBase, offset=0):
        if offset > 0:
            stream.seek(offset)

        self.signature = unpack_type(UINT, stream.read(4))
        if self.signature != PE_MAGIC and self.signature != PE_CIGAM:
            raise TypeError("Not a valid PE (Invalid NT signature)")

        self.file_header = COFFHeader(stream, self.endianness)

        # Parse optional header
        self.optional_header = None
        jmp_op = stream.tell()
        if self.file_header.size_opheader > 0:
            self.optional_header = OptionalHeader(stream, self.endianness)
        stream.seek(jmp_op + self.file_header.size_opheader)

    def __str__(self):
        return str(self.file_header)

    @property
    def endianness(self):
        return "big" if self.signature == PE_MAGIC else "little"

class Pe(object):
    def __init__(self, stream: io.RawIOBase):
        self.dos_header = DosHeader(stream)
        self.pe_header = PEHeader(stream, self.dos_header.e_lfanew)

        # Parse sections
        self.sections = []
        for _ in range(self.pe_header.file_header.nsections):
            self.sections.append(PeSectionHeader(stream, self.pe_header.endianness))

    def __str__(self):
        return "\n".join(["Pe Header",
                          "Magic:                   0x%02x",
                          "%s"]) % (self.pe_header.signature, str(self.pe_header))

    @staticmethod
    def verify(file: io.RawIOBase):
        s_pos = file.tell()
        magic = unpack_type(USHORT, file.read(2))  # MS_WORD
        file.seek(s_pos)
        return magic == MZ_MAGIC or magic == MZ_CIGAM


from uuid import uuid1

import click

from awsscripter.stack.helpers import catch_exceptions, confirmation
from awsscripter.stack.helpers import simplify_change_set_description
from awsscripter.stack.helpers import write, get_stack_or_env
from awsscripter.stack.stack_status import StackStatus, StackChangeSetStatus

@click.command(name="update")
@click.argument("path")
@click.option(
    "-c", "--change-set", is_flag=True,
    help="Create a change set before updating."
)
@click.option(
    "-v", "--verbose", is_flag=True, help="Display verbose output."
)
@click.option(
    "-y", "--yes", is_flag=True, help="Assume yes to all questions."
)
@click.pass_context
@catch_exceptions
def update_command(ctx, path, change_set, verbose, yes):
    """
    Update a stack.

    Updates a stack for a given config PATH. Or perform an update via
    change-set when the change-set flag is set.
    """

    stack, _ = get_stack_or_env(ctx, path)
    if change_set:
        change_set_name = "-".join(["change-set", uuid1().hex])
        stack.create_change_set(change_set_name)
        try:
            # Wait for change set to be created
            status = stack.wait_for_cs_completion(change_set_name)

            # Exit if change set fails to create
            if status != StackChangeSetStatus.READY:
                exit(1)

            # Describe changes
            description = stack.describe_change_set(change_set_name)
            if not verbose:
                description = simplify_change_set_description(description)
            write(description, ctx.obj["output_format"])

            # Execute change set if happy with changes
            if yes or click.confirm("Proceed with stack update?"):
                stack.execute_change_set(change_set_name)
        finally:
            # Clean up by deleting change set
            stack.delete_change_set(change_set_name)
    else:
        confirmation("update", yes, stack=path)
        response = stack.update()
        if response != StackStatus.COMPLETE:
            exit(1)


#
# License: See LICENSE.md file
# GitHub: https://github.com/Baekalfen/PyBoy
#
__pdoc__ = {
    "GameWrapperPokemonGen1.cartridge_title": False,
    "GameWrapperPokemonGen1.post_tick": False,
}

import logging

from pyboy.utils import WindowEvent

from .base_plugin import PyBoyGameWrapper

logger = logging.getLogger(__name__)

try:
    from cython import compiled
    cythonmode = compiled
except ImportError:
    cythonmode = False

class GameWrapperPokemonGen1(PyBoyGameWrapper):
    """
    This class wraps Pokemon Red/Blue, and provides basic access for AIs.

    If you call `print` on an instance of this object, it will show an overview of everything this object provides.
    """
    cartridge_title = None

    def __init__(self, *args, **kwargs):
        self.shape = (20, 18)
        super().__init__(*args, game_area_section=(0, 0) + self.shape, game_area_wrap_around=True, **kwargs)
        self.sprite_offset = 0x1000

    def enabled(self):
        return self.pyboy_argv.get("game_wrapper") and ((self.pyboy.cartridge_title() == "POKEMON RED") or
                                                        (self.pyboy.cartridge_title() == "POKEMON BLUE"))

    def post_tick(self):
        self._tile_cache_invalid = True
        self._sprite_cache_invalid = True

        scanline_parameters = self.pyboy.botsupport_manager().screen().tilemap_position_list()
        WX = scanline_parameters[0][2]
        WY = scanline_parameters[0][3]
        self.use_background(WY != 0)

    def __repr__(self):
        adjust = 4
        # yapf: disable
        return (
            f"Pokemon Gen 1:\n" +
            "Sprites on screen:\n" +
            "\n".join([str(s) for s in self._sprites_on_screen()]) +
            "\n" +
            "Tiles on screen:\n" +
            " "*5 + "".join([f"{i: <4}" for i in range(10)]) + "\n" +
            "_"*(adjust*20+4) +
            "\n" +
            "\n".join(
                [
                    f"{i: <3}| " + "".join([str(tile).ljust(adjust) for tile in line])
                    for i, line in enumerate(self.game_area())
                ]
            )
        )
        # yapf: enable