# -*- coding: utf-8 -*-
"""The :class:`Schema` class, including its metaclass and options (class Meta)."""
from __future__ import absolute_import, unicode_literals

from collections import defaultdict, Mapping, OrderedDict
import functools
import copy
import datetime as dt
import decimal
import inspect
import json
import uuid
import warnings

from marshmallow import base, fields, utils, class_registry, marshalling
from marshmallow.compat import (
    binary_type,
    iteritems,
    iterkeys,
    text_type,
    with_metaclass,
)
from marshmallow.exceptions import ValidationError, StringNotCollectionError
from marshmallow.orderedset import OrderedSet
from marshmallow.decorators import (
    POST_DUMP,
    POST_LOAD,
    PRE_DUMP,
    PRE_LOAD,
    VALIDATES,
    VALIDATES_SCHEMA,
)
from marshmallow.utils import RAISE, missing, is_collection

def _get_fields(attrs, field_class, pop=False, ordered=False):
    """Get fields from a class. If ordered=True, fields will sorted by creation index.

    :param attrs: Mapping of class attributes
    :param type field_class: Base field class
    :param bool pop: Remove matching fields
    """
    fields = [
        (field_name, field_value)
        for field_name, field_value in iteritems(attrs)
        if utils.is_instance_or_subclass(field_value, field_class)
    ]
    if pop:
        for field_name, _ in fields:
            del attrs[field_name]
    if ordered:
        fields.sort(key=lambda pair: pair[1]._creation_index)
    return fields

# This function allows Schemas to inherit from non-Schema classes and ensures
#   inheritance according to the MRO
def _get_fields_by_mro(klass, field_class, ordered=False):
    """Collect fields from a class, following its method resolution order. The
    class itself is excluded from the search; only its parents are checked. Get
    fields from ``_declared_fields`` if available, else use ``__dict__``.

    :param type klass: Class whose fields to retrieve
    :param type field_class: Base field class
    """
    mro = inspect.getmro(klass)
    # Loop over mro in reverse to maintain correct order of fields
    return sum(
        (
            _get_fields(
                getattr(base, '_declared_fields', base.__dict__),
                field_class,
                ordered=ordered,
            )
            for base in mro[:0:-1]
        ),
        [],
    )

class SchemaMeta(type):
    """Metaclass for the Schema class. Binds the declared fields to
    a ``_declared_fields`` attribute, which is a dictionary mapping attribute
    names to field objects. Also sets the ``opts`` class attribute, which is
    the Schema class's ``class Meta`` options.
    """

    def __new__(mcs, name, bases, attrs):
        meta = attrs.get('Meta')
        ordered = getattr(meta, 'ordered', False)
        if not ordered:
            # Inherit 'ordered' option
            # Warning: We loop through bases instead of MRO because we don't
            # yet have access to the class object
            # (i.e. can't call super before we have fields)
            for base_ in bases:
                if hasattr(base_, 'Meta') and hasattr(base_.Meta, 'ordered'):
                    ordered = base_.Meta.ordered
                    break
            else:
                ordered = False
        cls_fields = _get_fields(attrs, base.FieldABC, pop=True, ordered=ordered)
        klass = super(SchemaMeta, mcs).__new__(mcs, name, bases, attrs)
        inherited_fields = _get_fields_by_mro(klass, base.FieldABC, ordered=ordered)

        # Use getattr rather than attrs['Meta'] so that we get inheritance for free
        meta = getattr(klass, 'Meta')
        # Set klass.opts in __new__ rather than __init__ so that it is accessible in
        # get_declared_fields
        klass.opts = klass.OPTIONS_CLASS(meta, ordered=ordered)
        # Add fields specifid in the `include` class Meta option
        cls_fields += list(klass.opts.include.items())

        dict_cls = OrderedDict if ordered else dict
        # Assign _declared_fields on class
        klass._declared_fields = mcs.get_declared_fields(
            klass=klass,
            cls_fields=cls_fields,
            inherited_fields=inherited_fields,
            dict_cls=dict_cls,
        )
        return klass

    @classmethod
    def get_declared_fields(mcs, klass, cls_fields, inherited_fields, dict_cls):
        """Returns a dictionary of field_name => `Field` pairs declard on the class.
        This is exposed mainly so that plugins can add additional fields, e.g. fields
        computed from class Meta options.

        :param type klass: The class object.
        :param dict cls_fields: The fields declared on the class, including those added
            by the ``include`` class Meta option.
        :param dict inherited_fileds: Inherited fields.
        :param type dict_class: Either `dict` or `OrderedDict`, depending on the whether
            the user specified `ordered=True`.
        """
        return dict_cls(inherited_fields + cls_fields)

    # NOTE: self is the class object
    def __init__(self, name, bases, attrs):
        super(SchemaMeta, self).__init__(name, bases, attrs)
        if name:
            class_registry.register(name, self)
        self._hooks = self.resolve_hooks()

    def resolve_hooks(self):
        """Add in the decorated processors

        By doing this after constructing the class, we let standard inheritance
        do all the hard work.
        """
        mro = inspect.getmro(self)

        hooks = defaultdict(list)

        for attr_name in dir(self):
            # Need to look up the actual descriptor, not whatever might be
            # bound to the class. This needs to come from the __dict__ of the
            # declaring class.
            for parent in mro:
                try:
                    attr = parent.__dict__[attr_name]
                except KeyError:
                    continue
                else:
                    break
            else:
                # In case we didn't find the attribute and didn't break above.
                # We should never hit this - it's just here for completeness
                # to exclude the possibility of attr being undefined.
                continue

            try:
                hook_config = attr.__marshmallow_hook__
            except AttributeError:
                pass
            else:
                for key in iterkeys(hook_config):
                    # Use name here so we can get the bound method later, in
                    # case the processor was a descriptor or something.
                    hooks[key].append(attr_name)

        return hooks

class SchemaOpts(object):
    """class Meta options for the :class:`Schema`. Defines defaults."""

    def __init__(self, meta, ordered=False):
        self.fields = getattr(meta, 'fields', ())
        if not isinstance(self.fields, (list, tuple)):
            raise ValueError('`fields` option must be a list or tuple.')
        self.additional = getattr(meta, 'additional', ())
        if not isinstance(self.additional, (list, tuple)):
            raise ValueError('`additional` option must be a list or tuple.')
        if self.fields and self.additional:
            raise ValueError(
                'Cannot set both `fields` and `additional` options'
                ' for the same Schema.',
            )
        self.exclude = getattr(meta, 'exclude', ())
        if not isinstance(self.exclude, (list, tuple)):
            raise ValueError('`exclude` must be a list or tuple.')
        self.dateformat = getattr(meta, 'dateformat', None)
        if hasattr(meta, 'json_module'):
            warnings.warn(
                'The json_module class Meta option is deprecated. Use render_module instead.',
                DeprecationWarning,
            )
            render_module = getattr(meta, 'json_module', json)
        else:
            render_module = json
        self.render_module = getattr(meta, 'render_module', render_module)
        self.ordered = getattr(meta, 'ordered', ordered)
        self.index_errors = getattr(meta, 'index_errors', True)
        self.include = getattr(meta, 'include', {})
        self.load_only = getattr(meta, 'load_only', ())
        self.dump_only = getattr(meta, 'dump_only', ())
        self.unknown = getattr(meta, 'unknown', RAISE)

class BaseSchema(base.SchemaABC):
    """Base schema class with which to define custom schemas.

    Example usage:

    .. code-block:: python

        import datetime as dt
        from marshmallow import Schema, fields

        class Album(object):
            def __init__(self, title, release_date):
                self.title = title
                self.release_date = release_date

        class AlbumSchema(Schema):
            title = fields.Str()
            release_date = fields.Date()

        # Or, equivalently
        class AlbumSchema2(Schema):
            class Meta:
                fields = ("title", "release_date")

        album = Album("Beggars Banquet", dt.date(1968, 12, 6))
        schema = AlbumSchema()
        data, errors = schema.dump(album)
        data  # {'release_date': '1968-12-06', 'title': 'Beggars Banquet'}

    :param tuple|list only: Whitelist of the declared fields to select when
        instantiating the Schema. If None, all fields are used. Nested fields
        can be represented with dot delimiters.
    :param tuple|list exclude: Blacklist of the declared fields to exclude
        when instantiating the Schema. If a field appears in both `only` and
        `exclude`, it is not used. Nested fields can be represented with dot
        delimiters.
    :param str prefix: Optional prefix that will be prepended to all the
        serialized field names.
    :param bool many: Should be set to `True` if ``obj`` is a collection
        so that the object will be serialized to a list.
    :param dict context: Optional context passed to :class:`fields.Method` and
        :class:`fields.Function` fields.
    :param tuple|list load_only: Fields to skip during serialization (write-only fields)
    :param tuple|list dump_only: Fields to skip during deserialization (read-only fields)
    :param bool|tuple partial: Whether to ignore missing fields. If its value
        is an iterable, only missing fields listed in that iterable will be
        ignored.
    :param unknown: Whether to exclude, include, or raise an error for unknown
        fields in the data. Use `EXCLUDE`, `INCLUDE` or `RAISE`.

    .. versionchanged:: 2.0.0
        `__validators__`, `__preprocessors__`, and `__data_handlers__` are removed in favor of
        `marshmallow.decorators.validates_schema`,
        `marshmallow.decorators.pre_load` and `marshmallow.decorators.post_dump`.
        `__accessor__` and `__error_handler__` are deprecated. Implement the
        `handle_error` and `get_attribute` methods instead.
        """
    TYPE_MAPPING = {
        text_type: fields.String,
        binary_type: fields.String,
        dt.datetime: fields.DateTime,
        float: fields.Float,
        bool: fields.Boolean,
        tuple: fields.Raw,
        list: fields.Raw,
        set: fields.Raw,
        int: fields.Integer,
        uuid.UUID: fields.UUID,
        dt.time: fields.Time,
        dt.date: fields.Date,
        dt.timedelta: fields.TimeDelta,
        decimal.Decimal: fields.Decimal,
    }

    OPTIONS_CLASS = SchemaOpts

    class Meta(object):
        """Options object for a Schema.

        Example usage: ::

            class Meta:
                fields = ("id", "email", "date_created")
                exclude = ("password", "secret_attribute")

        Available options:

        - ``fields``: Tuple or list of fields to include in the serialized result.
        - ``additional``: Tuple or list of fields to include *in addition* to the
            explicitly declared fields. ``additional`` and ``fields`` are
            mutually-exclusive options.
        - ``include``: Dictionary of additional fields to include in the schema. It is
            usually better to define fields as class variables, but you may need to
            use this option, e.g., if your fields are Python keywords. May be an
            `OrderedDict`.
        - ``exclude``: Tuple or list of fields to exclude in the serialized result.
            Nested fields can be represented with dot delimiters.
        - ``dateformat``: Date format for all DateTime fields that do not have their
            date format explicitly specified.
        - ``render_module``: Module to use for `loads` and `dumps`. Defaults to
            `json` from the standard library.
            Defaults to the ``json`` module in the stdlib.
        - ``ordered``: If `True`, order serialization output according to the
            order in which fields were declared. Output of `Schema.dump` will be a
            `collections.OrderedDict`.
        - ``index_errors``: If `True`, errors dictionaries will include the index
            of invalid items in a collection.
        - ``load_only``: Tuple or list of fields to exclude from serialized results.
        - ``dump_only``: Tuple or list of fields to exclude from deserialization
        - ``unknown``: Whether to exclude, include, or raise an error for unknown
            fields in the data. Use `EXCLUDE`, `INCLUDE` or `RAISE`.
        """
        pass

    def __init__(
        self, only=None, exclude=(), prefix='', many=False,
        context=None, load_only=(), dump_only=(), partial=False,
        unknown=None,
    ):
        # Raise error if only or exclude is passed as string, not list of strings
        if only is not None and not is_collection(only):
            raise StringNotCollectionError('"only" should be a list of strings')
        if exclude is not None and not is_collection(exclude):
            raise StringNotCollectionError('"exclude" should be a list of strings')
        # copy declared fields from metaclass
        self.declared_fields = copy.deepcopy(self._declared_fields)
        self.many = many
        self.only = only
        self.exclude = exclude
        self.prefix = prefix
        self.ordered = self.opts.ordered
        self.load_only = set(load_only) or set(self.opts.load_only)
        self.dump_only = set(dump_only) or set(self.opts.dump_only)
        self.partial = partial
        self.unknown = unknown or self.opts.unknown
        #: Dictionary mapping field_names -> :class:`Field` objects
        self.fields = self.dict_class()
        self.context = context or {}
        self._normalize_nested_options()
        self._types_seen = set()
        self._update_fields(many=many)

    def __repr__(self):
        return '<{ClassName}(many={self.many})>'.format(
            ClassName=self.__class__.__name__, self=self,
        )

    @property
    def dict_class(self):
        return OrderedDict if self.ordered else dict

    @property
    def set_class(self):
        return OrderedSet if self.ordered else set

    ##### Override-able methods #####

    def handle_error(self, error, data):
        """Custom error handler function for the schema.

        :param ValidationError error: The `ValidationError` raised during (de)serialization.
        :param data: The original input data.

        .. versionadded:: 2.0.0
        """
        pass

    def get_attribute(self, obj, attr, default):
        """Defines how to pull values from an object to serialize.

        .. versionadded:: 2.0.0

        .. versionchanged:: 3.0.0a1
            Changed position of ``obj`` and ``attr``.
        """
        return utils.get_value(obj, attr, default)

    ##### Serialization/Deserialization API #####

    def dump(self, obj, many=None, update_fields=True):
        """Serialize an object to native Python data types according to this
        Schema's fields.

        :param obj: The object to serialize.
        :param bool many: Whether to serialize `obj` as a collection. If `None`, the value
            for `self.many` is used.
        :param bool update_fields: Whether to update the schema's field classes. Typically
            set to `True`, but may be `False` when serializing a homogenous collection.
            This parameter is used by `fields.Nested` to avoid multiple updates.
        :return: A dict of serialized data
        :rtype: dict

        .. versionadded:: 1.0.0
        .. versionchanged:: 3.0.0b7
            This method returns the serialized data rather than a ``(data, errors)`` duple.
            A :exc:`ValidationError <marshmallow.exceptions.ValidationError>` is raised
            if ``obj`` is invalid.
        """
        # Callable marshalling object
        marshal = marshalling.Marshaller(prefix=self.prefix)
        errors = {}
        many = self.many if many is None else bool(many)
        if many and utils.is_iterable_but_not_string(obj):
            obj = list(obj)

        if self._has_processors(PRE_DUMP):
            try:
                processed_obj = self._invoke_dump_processors(
                    PRE_DUMP,
                    obj,
                    many,
                    original_data=obj,
                )
            except ValidationError as error:
                errors = error.normalized_messages()
                result = None
        else:
            processed_obj = obj

        if not errors:
            if update_fields:
                obj_type = type(processed_obj)
                if obj_type not in self._types_seen:
                    self._update_fields(processed_obj, many=many)
                    if not isinstance(processed_obj, Mapping):
                        self._types_seen.add(obj_type)

            try:
                result = marshal(
                    processed_obj,
                    self.fields,
                    many=many,
                    accessor=self.get_attribute,
                    dict_class=self.dict_class,
                    index_errors=self.opts.index_errors,
                )
            except ValidationError as error:
                errors = marshal.errors
                result = error.data

        if not errors and self._has_processors(POST_DUMP):
            try:
                result = self._invoke_dump_processors(
                    POST_DUMP,
                    result,
                    many,
                    original_data=obj,
                )
            except ValidationError as error:
                errors = error.normalized_messages()
        if errors:
            exc = ValidationError(
                errors,
                field_names=marshal.error_field_names,
                data=obj,
                valid_data=result,
                **marshal.error_kwargs
            )
            # User-defined error handler
            self.handle_error(exc, obj)
            raise exc

        return result

    def dumps(self, obj, many=None, update_fields=True, *args, **kwargs):
        """Same as :meth:`dump`, except return a JSON-encoded string.

        :param obj: The object to serialize.
        :param bool many: Whether to serialize `obj` as a collection. If `None`, the value
            for `self.many` is used.
        :param bool update_fields: Whether to update the schema's field classes. Typically
            set to `True`, but may be `False` when serializing a homogenous collection.
            This parameter is used by `fields.Nested` to avoid multiple updates.
        :return: A ``json`` string
        :rtype: str

        .. versionadded:: 1.0.0
        .. versionchanged:: 3.0.0b7
            This method returns the serialized data rather than a ``(data, errors)`` duple.
            A :exc:`ValidationError <marshmallow.exceptions.ValidationError>` is raised
            if ``obj`` is invalid.
        """
        serialized = self.dump(obj, many=many, update_fields=update_fields)
        return self.opts.render_module.dumps(serialized, *args, **kwargs)

    def load(self, data, many=None, partial=None, unknown=None):
        """Deserialize a data structure to an object defined by this Schema's
        fields and :meth:`make_object`.

        :param dict data: The data to deserialize.
        :param bool many: Whether to deserialize `data` as a collection. If `None`, the
            value for `self.many` is used.
        :param bool|tuple partial: Whether to ignore missing fields. If `None`,
            the value for `self.partial` is used. If its value is an iterable,
            only missing fields listed in that iterable will be ignored.
        :param unknown: Whether to exclude, include, or raise an error for unknown
            fields in the data. Use `EXCLUDE`, `INCLUDE` or `RAISE`.
            If `None`, the value for `self.unknown` is used.
        :return: A dict of deserialized data
        :rtype: dict

        .. versionadded:: 1.0.0
        .. versionchanged:: 3.0.0b7
            This method returns the deserialized data rather than a ``(data, errors)`` duple.
            A :exc:`ValidationError <marshmallow.exceptions.ValidationError>` is raised
            if invalid data are passed.
        """
        return self._do_load(
            data, many, partial=partial, unknown=unknown,
            postprocess=True,
        )

    def loads(
        self, json_data, many=None, partial=None, unknown=None,
        **kwargs
    ):
        """Same as :meth:`load`, except it takes a JSON string as input.

        :param str json_data: A JSON string of the data to deserialize.
        :param bool many: Whether to deserialize `obj` as a collection. If `None`, the
            value for `self.many` is used.
        :param bool|tuple partial: Whether to ignore missing fields. If `None`,
            the value for `self.partial` is used. If its value is an iterable,
            only missing fields listed in that iterable will be ignored.
        :param unknown: Whether to exclude, include, or raise an error for unknown
            fields in the data. Use `EXCLUDE`, `INCLUDE` or `RAISE`.
            If `None`, the value for `self.unknown` is used.
        :return: A dict of deserialized data
        :rtype: dict

        .. versionadded:: 1.0.0
        .. versionchanged:: 3.0.0b7
            This method returns the deserialized data rather than a ``(data, errors)`` duple.
            A :exc:`ValidationError <marshmallow.exceptions.ValidationError>` is raised
            if invalid data are passed.
        """
        data = self.opts.render_module.loads(json_data, **kwargs)
        return self.load(data, many=many, partial=partial, unknown=unknown)

    def validate(self, data, many=None, partial=None):
        """Validate `data` against the schema, returning a dictionary of
        validation errors.

        :param dict data: The data to validate.
        :param bool many: Whether to validate `data` as a collection. If `None`, the
            value for `self.many` is used.
        :param bool|tuple partial: Whether to ignore missing fields. If `None`,
            the value for `self.partial` is used. If its value is an iterable,
            only missing fields listed in that iterable will be ignored.
        :return: A dictionary of validation errors.
        :rtype: dict

        .. versionadded:: 1.1.0
        """
        try:
            self._do_load(data, many, partial=partial, postprocess=False)
        except ValidationError as exc:
            return exc.messages
        return {}

    ##### Private Helpers #####

    def _do_load(
        self, data, many=None, partial=None, unknown=None,
        postprocess=True,
    ):
        """Deserialize `data`, returning the deserialized result.

        :param data: The data to deserialize.
        :param bool many: Whether to deserialize `data` as a collection. If `None`, the
            value for `self.many` is used.
        :param bool|tuple partial: Whether to validate required fields. If its value is an iterable,
            only fields listed in that iterable will be ignored will be allowed missing.
            If `True`, all fields will be allowed missing.
            If `None`, the value for `self.partial` is used.
        :param unknown: Whether to exclude, include, or raise an error for unknown
            fields in the data. Use `EXCLUDE`, `INCLUDE` or `RAISE`.
            If `None`, the value for `self.unknown` is used.
        :param bool postprocess: Whether to run post_load methods..
        :return: A dict of deserialized data
        :rtype: dict
        """
        # Callable unmarshalling object
        unmarshal = marshalling.Unmarshaller()
        errors = {}
        many = self.many if many is None else bool(many)
        unknown = unknown or self.unknown
        if partial is None:
            partial = self.partial
        if self._has_processors(PRE_LOAD):
            try:
                processed_data = self._invoke_load_processors(
                    PRE_LOAD,
                    data,
                    many,
                    original_data=data,
                )
            except ValidationError as err:
                errors = err.normalized_messages()
                result = None
        else:
            processed_data = data
        if not errors:
            try:
                result = unmarshal(
                    processed_data,
                    self.fields,
                    many=many,
                    partial=partial,
                    unknown=unknown,
                    dict_class=self.dict_class,
                    index_errors=self.opts.index_errors,
                )
            except ValidationError as error:
                result = error.data
            self._invoke_field_validators(unmarshal, data=result, many=many)
            errors = unmarshal.errors
            # Run schema-level validation.
            if self._has_processors(VALIDATES_SCHEMA):
                field_errors = bool(errors)
                try:
                    self._invoke_schema_validators(
                        unmarshal,
                        pass_many=True,
                        data=result,
                        original_data=data,
                        many=many,
                        field_errors=field_errors,
                    )
                except ValidationError as err:
                    errors.update(err.messages)
                try:
                    self._invoke_schema_validators(
                        unmarshal,
                        pass_many=False,
                        data=result,
                        original_data=data,
                        many=many,
                        field_errors=field_errors,
                    )
                except ValidationError as err:
                    errors.update(err.messages)
        # Run post processors
        if not errors and postprocess and self._has_processors(POST_LOAD):
            try:
                result = self._invoke_load_processors(
                    POST_LOAD,
                    result,
                    many,
                    original_data=data,
                )
            except ValidationError as err:
                errors = err.normalized_messages()
        if errors:
            exc = ValidationError(
                errors,
                field_names=unmarshal.error_field_names,
                data=data,
                valid_data=result,
                **unmarshal.error_kwargs
            )
            self.handle_error(exc, data)
            raise exc

        return result

    def _normalize_nested_options(self):
        """Apply then flatten nested schema options"""
        if self.only is not None:
            # Apply the only option to nested fields.
            self.__apply_nested_option('only', self.only, 'intersection')
            # Remove the child field names from the only option.
            self.only = self.set_class(
                [field.split('.', 1)[0] for field in self.only],
            )
        excludes = set(self.opts.exclude) | set(self.exclude)
        if excludes:
            # Apply the exclude option to nested fields.
            self.__apply_nested_option('exclude', excludes, 'union')
        if self.exclude:
            # Remove the parent field names from the exclude option.
            self.exclude = self.set_class(
                [field for field in self.exclude if '.' not in field],
            )
        if self.opts.exclude:
            # Remove the parent field names from the meta exclude option.
            self.opts.exclude = self.set_class(
                [field for field in self.opts.exclude if '.' not in field],
            )

    def __apply_nested_option(self, option_name, field_names, set_operation):
        """Apply nested options to nested fields"""
        # Split nested field names on the first dot.
        nested_fields = [name.split('.', 1) for name in field_names if '.' in name]
        # Partition the nested field names by parent field.
        nested_options = defaultdict(list)
        for parent, nested_names in nested_fields:
            nested_options[parent].append(nested_names)
        # Apply the nested field options.
        for key, options in iter(nested_options.items()):
            new_options = self.set_class(options)
            original_options = getattr(self.declared_fields[key], option_name, ())
            if original_options:
                if set_operation == 'union':
                    new_options |= self.set_class(original_options)
                if set_operation == 'intersection':
                        new_options &= self.set_class(original_options)
            setattr(self.declared_fields[key], option_name, new_options)

    def _update_fields(self, obj=None, many=False):
        """Update fields based on the passed in object."""
        if self.opts.fields:
            # Return fields specified in fields option
            field_names = self.set_class(self.opts.fields)
        elif self.opts.additional:
            # Return declared fields + additional fields
            field_names = (self.set_class(self.declared_fields.keys()) |
                            self.set_class(self.opts.additional))
        else:
            field_names = self.set_class(self.declared_fields.keys())

        declared_fields = field_names
        invalid_fields = self.set_class()
        if self.only is not None:
            # Return only fields specified in only option
            only = self.set_class(self.only)
            field_names = only
            invalid_only = only - declared_fields
            invalid_fields |= invalid_only

        # If "exclude" option or param is specified, remove those fields
        excludes = set(self.opts.exclude) | set(self.exclude)
        if excludes:
            field_names = field_names - excludes
            invalid_excludes = excludes - declared_fields
            invalid_fields |= invalid_excludes

        if invalid_fields:
            message = 'Invalid fields for {0}: {1}.'.format(self, invalid_fields)
            raise ValueError(message)

        ret = self.__filter_fields(field_names, obj, many=many)
        # Set parents
        self.__set_field_attrs(ret)
        self.fields = ret
        return self.fields

    def on_bind_field(self, field_name, field_obj):
        """Hook to modify a field when it is bound to the `Schema`. No-op by default."""
        return None

    def __set_field_attrs(self, fields_dict):
        """Bind fields to the schema, setting any necessary attributes
        on the fields (e.g. parent and name).

        Also set field load_only and dump_only values if field_name was
        specified in ``class Meta``.
        """
        for field_name, field_obj in iteritems(fields_dict):
            try:
                if field_name in self.load_only:
                    field_obj.load_only = True
                if field_name in self.dump_only:
                    field_obj.dump_only = True
                field_obj._add_to_schema(field_name, self)
                self.on_bind_field(field_name, field_obj)
            except TypeError:
                # field declared as a class, not an instance
                if (isinstance(field_obj, type) and
                        issubclass(field_obj, base.FieldABC)):
                    msg = ('Field for "{0}" must be declared as a '
                           'Field instance, not a class. '
                           'Did you mean "fields.{1}()"?'
                           .format(field_name, field_obj.__name__))
                    raise TypeError(msg)
        return fields_dict

    def __filter_fields(self, field_names, obj, many=False):
        """Return only those field_name:field_obj pairs specified by
        ``field_names``.

        :param set field_names: Field names to include in the final
            return dictionary.
        :returns: An dict of field_name:field_obj pairs.
        """
        if obj and many:
            try:  # Homogeneous collection
                # Prefer getitem over iter to prevent breaking serialization
                # of objects for which iter will modify position in the collection
                # e.g. Pymongo cursors
                if hasattr(obj, '__getitem__') and callable(getattr(obj, '__getitem__')):
                    try:
                        obj_prototype = obj[0]
                    except KeyError:
                        obj_prototype = next(iter(obj))
                else:
                    obj_prototype = next(iter(obj))
            except (StopIteration, IndexError):  # Nothing to serialize
                return {k: v for k, v in self.declared_fields.items() if k in field_names}
            obj = obj_prototype
        ret = self.dict_class()
        for key in field_names:
            if key in self.declared_fields:
                ret[key] = self.declared_fields[key]
            else:  # Implicit field creation (class Meta 'fields' or 'additional')
                if obj:
                    attribute_type = None
                    try:
                        if isinstance(obj, Mapping):
                            attribute_type = type(obj[key])
                        else:
                            attribute_type = type(getattr(obj, key))
                    except (AttributeError, KeyError) as err:
                        err_type = type(err)
                        raise err_type(
                            '"{0}" is not a valid field for {1}.'.format(key, obj),
                        )
                    field_obj = self.TYPE_MAPPING.get(attribute_type, fields.Field)()
                else:  # Object is None
                    field_obj = fields.Field()
                # map key -> field (default to Raw)
                ret[key] = field_obj
        return ret

    def _has_processors(self, tag):
        return self._hooks[(tag, True)] or self._hooks[(tag, False)]

    def _invoke_dump_processors(self, tag, data, many, original_data=None):
        # The pass_many post-dump processors may do things like add an envelope, so
        # invoke those after invoking the non-pass_many processors which will expect
        # to get a list of items.
        data = self._invoke_processors(
            tag, pass_many=False,
            data=data, many=many, original_data=original_data,
        )
        data = self._invoke_processors(
            tag, pass_many=True,
            data=data, many=many, original_data=original_data,
        )
        return data

    def _invoke_load_processors(self, tag, data, many, original_data=None):
        # This has to invert the order of the dump processors, so run the pass_many
        # processors first.
        data = self._invoke_processors(
            tag, pass_many=True,
            data=data, many=many, original_data=original_data,
        )
        data = self._invoke_processors(
            tag, pass_many=False,
            data=data, many=many, original_data=original_data,
        )
        return data

    def _invoke_field_validators(self, unmarshal, data, many):
        for attr_name in self._hooks[VALIDATES]:
            validator = getattr(self, attr_name)
            validator_kwargs = validator.__marshmallow_hook__[VALIDATES]
            field_name = validator_kwargs['field_name']

            try:
                field_obj = self.fields[field_name]
            except KeyError:
                if field_name in self.declared_fields:
                    continue
                raise ValueError('"{0}" field does not exist.'.format(field_name))

            if many:
                for idx, item in enumerate(data):
                    try:
                        value = item[field_obj.attribute or field_name]
                    except KeyError:
                        pass
                    else:
                        validated_value = unmarshal.call_and_store(
                            getter_func=validator,
                            data=value,
                            field_name=field_obj.data_key or field_name,
                            index=(idx if self.opts.index_errors else None),
                        )
                        if validated_value is missing:
                            data[idx].pop(field_name, None)
            else:
                try:
                    value = data[field_obj.attribute or field_name]
                except KeyError:
                    pass
                else:
                    validated_value = unmarshal.call_and_store(
                        getter_func=validator,
                        data=value,
                        field_name=field_obj.data_key or field_name,
                    )
                    if validated_value is missing:
                        data.pop(field_name, None)

    def _invoke_schema_validators(
        self,
        unmarshal,
        pass_many,
        data,
        original_data,
        many,
        field_errors=False,
    ):
        errors = {}
        for attr_name in self._hooks[(VALIDATES_SCHEMA, pass_many)]:
            validator = getattr(self, attr_name)
            validator_kwargs = validator.__marshmallow_hook__[(VALIDATES_SCHEMA, pass_many)]
            pass_original = validator_kwargs.get('pass_original', False)

            skip_on_field_errors = validator_kwargs['skip_on_field_errors']
            if skip_on_field_errors and field_errors:
                continue

            if pass_many:
                validator = functools.partial(validator, many=many)
            if many and not pass_many:
                for idx, (item, orig) in enumerate(zip(data, original_data)):
                    try:
                        unmarshal.run_validator(
                            validator,
                            item,
                            orig,
                            self.fields,
                            many=many,
                            index=idx,
                            pass_original=pass_original,
                        )
                    except ValidationError as err:
                        errors.update(err.messages)
            else:
                try:
                    unmarshal.run_validator(
                        validator,
                        data,
                        original_data,
                        self.fields,
                        many=many,
                        pass_original=pass_original,
                    )
                except ValidationError as err:
                    errors.update(err.messages)
        if errors:
            raise ValidationError(errors)
        return None

    def _invoke_processors(
        self,
        tag,
        pass_many,
        data,
        many,
        original_data=None,
    ):
        key = (tag, pass_many)
        for attr_name in self._hooks[key]:
            # This will be a bound method.
            processor = getattr(self, attr_name)

            processor_kwargs = processor.__marshmallow_hook__[key]
            pass_original = processor_kwargs.get('pass_original', False)

            if pass_many:
                if pass_original:
                    data = processor(data, many, original_data)
                else:
                    data = processor(data, many)
            elif many:
                if pass_original:
                    data = [
                        processor(item, original)
                        for item, original in zip(data, original_data)
                    ]
                else:
                    data = [processor(item) for item in data]
            else:
                if pass_original:
                    data = processor(data, original_data)
                else:
                    data = processor(data)
        return data

class Schema(with_metaclass(SchemaMeta, BaseSchema)):
    __doc__ = BaseSchema.__doc__


# /usr/bin/env python
# -*- encoding: utf-8 -*-
"""
【少前13-4自动刷图】
"""
# 打包命令pyinstaller main.spec
import logging.config
import os
import subprocess
import sys
import threading
import time
from functools import partial

import aircv
import win32api
import win32con
import yaml

# 全局覆盖print函数参数
print = partial(print, flush=True)

class Dnconsole:
    """
    【雷电控制台类】
    version: 9.0
    import该文件会自动实例化为 Dc
    """

    def __init__(self, installation_path: str):
        """
        【构造方法】
        """
        # if 模拟器安装路径存在性检测
        if os.path.exists(installation_path) is False:
            print('模拟器安装路径不存在！')
        # 获取模拟器安装路径
        self.ins_path = installation_path
        # Dnconsole程序路径
        self.console_path = self.ins_path + r'\ldconsole.exe '
        # ld程序路径
        self.ld_path = self.ins_path + r'\ld.exe '
        # if Dnconsole程序路径检测
        if os.path.exists(self.console_path) is False:
            print('Dnconsole程序路径不存在！\n请确认模拟器安装文件是否完整或模拟器版本是否不符！')
        # adb程序路径
        self.adb_path = self.ins_path + r'\adb.exe '
        # if adb程序路径检测
        if os.path.exists(self.adb_path) is False:
            print('Dnconsole程序路径不存在！\n请确认模拟器安装文件是否完整！')
        # 模拟器截屏程序路径
        self.screencap_path = r'/system/bin/screencap'
        # 模拟器截图保存路径
        self.devicess_path = r'/sdcard/Pictures/Screenshots/screenshot_tmp.png'
        # 用户环境
        self.userprofile = os.environ['USERPROFILE']
        self.documents_path = os.path.join(self.userprofile, 'Documents')
        # 本地图片保存路径
        self.images_path = self.documents_path + r'\leidian9\Pictures\Screenshots\\'
        # if 模拟器的截图保存路径
        if os.path.exists(self.images_path) is False:
            print('模拟器截图保存路径不存在！')
        # 获取当前工作目录路径
        self.workspace_path = os.getcwd()
        # 本地图片样本保存路径
        # self.target_path = self.workspace_path + r'\target\1080p_dpi280\\'
        self.target_path = self.workspace_path + r'\target\720p_dpi240\\'
        # 读取作战参数信息保存路径
        self.action_path = self.workspace_path + r'\script\\'
        # 读取作战参数
        # 构造完成
        logging.debug('Class-Dnconsole is ready.(%s)' % self.ins_path)
        # print('Class-Dnconsole is ready.(%s)' % self.ins_path)

    def YAML(self, script_name: str):
        """
        【读取Yaml文件】
        :param script_name:
        :return:
        """
        # 读取YAML参数
        yaml_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                                 self.action_path + script_name)
        logging.debug('yaml_path: %s' % yaml_path)
        try:
            # 打开文件
            with open(yaml_path, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
                return data
        except:
            return None

    def CMD(self, cmd: str):
        """
        【执行控制台命令语句】
        :param cmd: 命令
        :return: 控制台调试内容
        """
        CMD = self.console_path + cmd  # 控制台命令
        process = os.popen(CMD)
        result = process.read()
        process.close()
        return result

    def ldCMD(self, cmd: str):
        """
        【通过ld.exe执行代替adb的命令语句】
        :param cmd: 命令
        :return: 控制台调试内容
        """
        ldCMD = self.ld_path + cmd  # 控制台命令
        process = os.popen(ldCMD)
        result = process.read()
        process.close()
        return result

    def ADB(self, cmd: str):
        """
        【执行ADB命令语句】
        :param cmd: 命令
        :return: 控制台调试内容
        """
        CMD = self.adb_path + cmd  # adb命令
        # os.system('chcp 65001')
        p = subprocess.Popen(['C:\\leidian\\LDPlayer9\\adb', '-s', '127.0.0.1:5555', 'exec-out', 'screencap', '-p'],
                             stdout=subprocess.PIPE)
        output, _ = p.communicate()
        return output

    def launch(self, index: int = 0):
        """
        【启动模拟器】
        :param index: 模拟器序号
        :return: True=已启动 / False=不存在
        """
        cmd = 'launch --index %d' % index
        if self.CMD(cmd) == '':
            return True
        else:
            return False

    def runninglist(self):
        """
        【获取正在运行的模拟器列表（仅标题）】
        :return: 控制台调试内容
        """
        cmd = 'runninglist'
        return self.CMD(cmd)

    def isrunning(self, index: int = 0):
        """
        【检测模拟器是否启动】
        :param index: 模拟器序号
        :return: True=已启动 / False=未启动
        """
        cmd = 'isrunning --index %d' % index
        if self.CMD(cmd) == 'running':
            return True
        else:
            return False

    def runApp(self, index: int, packagename: str):
        """
        【运行App】
        :param index: 模拟器序号
        :param packagename: 包名
        :return: 控制台调试内容
        """
        cmd = 'runapp --index %d --packagename %s' % (index, packagename)
        return self.CMD(cmd)

    def isrunningAPP(self, index: int, packagename: str):
        """
        【检测应用是否启动】
        :param packagename: App包名
        :param index: 模拟器序号
        :return: True=已启动 / False=未启动
        """
        cmd = 'adb --index %d --command "shell dumpsys window | grep mCurrentFocus"' % index
        if packagename in self.CMD(cmd):
            return True
        else:
            return False

    def appIsrunning(self, index: int, packagename: str):
        """
        【获取App运行状态】
        :param index: 模拟器序号
        :param packagename: 包名
        :return: 控制台调试内容
        """
        cmd = 'adb --index %d --command "shell pidof %s"' % (index, packagename)
        return self.CMD(cmd)

    def killApp(self, index: int, packagename: str):
        """
        【终止App】
        :param index: 模拟器序号
        :param packagename: 包名
        :return: 控制台调试内容
        """
        cmd = 'killapp --index %d --packagename %s' % (index, packagename)
        return self.CMD(cmd)

    def screenShot(self, index: int):
        """
        【截屏并保存到本地】
        :param index: 模拟器序号
        """
        cmd1 = 'adb --index %d --command "shell %s -p %s"' % (index, self.screencap_path, self.devicess_path)
        cmd2 = 'adb --index %d --command "pull %s %s"' % (index, self.devicess_path, self.images_path)
        self.CMD(cmd1)
        self.CMD(cmd2)

    def screenShotnew(self, index: int):
        """
        【截屏并保存到本地】
        :param index: 模拟器序号
        """
        cmd = 'adb --index %d --command "shell %s -p %s"' % (index, self.screencap_path, self.devicess_path)
        result = self.CMD(cmd)
        return result

    def screenShotnewLd(self, index: int):
        """
        【截屏并保存到本地ld调用方式】
        :param index: 模拟器序号
        """
        cmd = '-s %d " %s -p %s"' % (index, self.screencap_path, self.devicess_path)
        result = self.ldCMD(cmd)
        scp_time = time.time() - start_time
        scp_time = time.strftime("%H:%M:%S", time.gmtime(scp_time))
        logging.info("%s 已截图" % scp_time)
        return result

    def actionOfTap(self, index: int, x: int, y: int):
        """
        【点击操作】
        :param index: 模拟器序号
        :param x: x
        :param y: y
        :return: 控制台调试内容
        """
        cmd = 'adb --index %d --command "shell input tap %d %d"' % (index, x, y)
        return self.CMD(cmd)

    def inputTap(self, index: int, x: int, y: int):
        """
        【点击操作】
        :param index: 模拟器序号
        :param x: x
        :param y: y
        :return: 控制台调试内容
        """
        cmd = '-s %d input tap %d %d' % (index, x, y)
        return self.ldCMD(cmd)

    def inputSwipe(self, index: int, x0: int, y0: int, x1: int, y1, ms: int = 200):
        """
        【滑动操作】
        :param index:
        :param x0:
        :param y0:
        :param x1:
        :param y1:
        :param ms:
        :return:
        """
        cmd = '-s %d input swipe %d %d %d %d %d' % (index, x0, y0, x1, y1, ms)
        return self.ldCMD(cmd)

class Ldaction(object):
    """
    【雷电操作类】
    import该文件需先实例化 Dconsole
    """

    def __init__(self, index: int, console, screenshot_img_name: str):
        """
        【构造方法】
        :param index: 模拟器号
        :param console: 传入的console对象
        :param screenshot_img_name:
        :return:
        """
        self.screenshot_img_file = console.images_path + screenshot_img_name
        # self.screenshot_img = aircv.imread(self.screenshot_img_file)
        self.index = index
        self.ld = console
        self.tap_info = self.ld.YAML

    def isExistV2(self, target_name: str, target_img_num: int = 0):
        """
        【判断指定图片是否包含在截图中】
        :param target_name:
        :param target_img_num:
        :return: 图片查询结果和包含详情
        """
        target_yaml = target_yaml_data[target_name]
        target_img_name = target_yaml['target_img_name_list'][target_img_num]
        rgb = target_yaml['rgb']
        threshold = target_yaml['threshold']
        screenshot_img = aircv.imread(self.screenshot_img_file)
        target_img_name_file = self.ld.target_path + target_img_name
        target_img = aircv.imread(target_img_name_file)
        result = aircv.find_template(screenshot_img, target_img, rgb=rgb, threshold=threshold)
        if result is None:
            return False, result
        else:
            return True, result

    def find_target_imgV4(self, target_name: str, tap_info: dict):
        """
        【改进后的目标图片搜索方法】
        :param tap_info:
        :param target_name:
        :return:
        """
        logging.debug('target: %s, 配置信息: %s' % (target_name, tap_info))
        target_img_name_list = tap_info['target_img_name_list']
        need_screenShot = tap_info['need_screenShot']
        moveX = tap_info['moveX']
        moveY = tap_info['moveY']
        tap_times = tap_info['tap_times']
        tap_interval = tap_info['tap_interval']
        search_again_times = tap_info['search_again_times']
        search_again_sleep_time = tap_info['search_again_sleep_time']
        for target_img_name in target_img_name_list:
            for search in range(search_again_times):
                if need_screenShot:
                    self.ld.ScreenShot_Ld(self.index)
                target_img_num = target_img_name_list.index(target_img_name)
                result = self.isExistV2(target_name, target_img_num)
                find_result = result[1]
                if find_result is None and need_screenShot is False and search_again_times == 1:
                    running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                    # print('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                    logging.error('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                    break
                elif find_result is None and search_again_times == 1:
                    running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                    # print('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                    logging.error('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                elif find_result is None and search_again_times > 1:
                    running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                    # print('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name), end="")
                    logging.warning('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                    search_times = 1
                    search_times += search
                    if search_again_sleep_time != 0:
                        search_again_sleep_time_double = int(search_again_sleep_time * 2)
                        for s in range(search_again_sleep_time_double + 1):
                            running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                            half_second = s
                            s_10, d = divmod(half_second, 10)
                            print("\r%s %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                            if s == search_again_sleep_time_double:
                                pass
                            else:
                                time.sleep(0.5)
                        print('\n', end='')
                    else:
                        pass
                    running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                    # print('%s 重新搜索 %s 按钮 %d 次' % (running_time, target_img_name, search_times))
                    logging.info('%s 重新搜索 %s 按钮 %d 次' % (running_time, target_img_name, search_times))
                    # 判断当不用截图的target节点当search_again_times大于1时重试时触发截图
                    if need_screenShot is False:
                        self.ld.ScreenShot_Ld(self.index)
                    if search_times == search_again_times:
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        # print('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                        logging.error('%s 未能找到需要的 %s 按钮！' % (running_time, target_img_name))
                else:
                    # 按钮坐标
                    x = find_result['result'][0]
                    y = find_result['result'][1]
                    # 点击坐标
                    tap_x = x - moveX
                    tap_y = y - moveY
                    for t in range(tap_times):
                        tap = t + 1
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        print('%s 找到按钮: %s [%d , %d]  ' % (running_time, target_img_name, x, y))
                        if tap_interval != 0:
                            tap_interval_time = int(tap_interval * 2)
                            for s in range(tap_interval_time + 1):
                                running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                                half_second = s
                                s_10, d = divmod(half_second, 10)
                                print("\r%s %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                                if s == tap_interval_time:
                                    pass
                                else:
                                    time.sleep(0.5)
                            print('\n', end='')
                        else:
                            pass
                        self.ld.inputTap(Console_Action.index, tap_x, tap_y)
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        print('%s 已点击按钮 %s [%d , %d] %d次' % (running_time, target_img_name, tap_x, tap_y, tap))
                    break
            if result[0]:
                break
        return result

    def LdactionTapV4(self, target_name: str, tap_info: dict):
        """
        【查找对应按钮】
        :param tap_info:
        :param target_name: 点击事件查询的按钮target
        """
        running_time = time.time() - start_time
        running_time = time.strftime("%H:%M:%S", time.gmtime(running_time))
        tap_result_check_name = tap_info['tap_result_check_name']
        before_tap_wait_time = tap_info['before_tap_wait_time']
        after_tap_wait_time = tap_info['after_tap_wait_time']
        beginning_content = tap_info['beginning_content']
        end_content = tap_info['end_content']
        error_content = tap_info['error_content']
        print('%s +++++START+++++ [目标%s: %s %s]' % (running_time, target_name, beginning_content, running_script))
        if before_tap_wait_time != 0:
            do_before_tap_wait_time = int(before_tap_wait_time * 2)
            for s in range(do_before_tap_wait_time + 1):
                running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                half_second = s
                s_10, d = divmod(half_second, 10)
                print("\r%s %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                if s == do_before_tap_wait_time:
                    pass
                else:
                    time.sleep(0.5)
            print('\n', end='')
        else:
            pass
        result = True
        while result:
            # 点击target
            find_result = self.find_target_imgV4(target_name, tap_info)
            # 判断点击方法处理结果
            if find_result[0]:
                # 处理点击后等待时间
                if after_tap_wait_time != 0:
                    do_after_tap_wait_time = int(after_tap_wait_time * 2)
                    for s in range(do_after_tap_wait_time + 1):
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        half_second = s
                        s_10, d = divmod(half_second, 10)
                        print("\r%s %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                        if s == do_after_tap_wait_time:
                            pass
                        else:
                            time.sleep(0.5)
                    print('\n', end='')
                else:
                    pass
                # 判断是否有校验对象
                if tap_result_check_name is not None:
                    # 赋值未校验标记
                    is_not_Exist_result = True
                    # 当有校验对象后执行循环校验
                    while is_not_Exist_result:
                        # 校验前截图
                        self.ld.ScreenShot_Ld(self.index)
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        # print('%s 开始校验 %s' % (running_time, tap_result_check_name))
                        logging.info('%s 开始校验 %s' % (running_time, tap_result_check_name))
                        # 执行校验
                        tap_result = self.isExistV2(tap_result_check_name)
                        # 判断校验结果，通过后结束校验
                        if tap_result[0]:
                            running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                            # print('%s 校验通过' % running_time)
                            logging.info('%s 校验通过' % running_time)
                            is_not_Exist_result = False
                            break
                        # 不通过重新点击target
                        elif tap_result[0] is False:
                            running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                            # print('%s 点击后结果校验不通过，重新点击！' % running_time)
                            logging.error('%s 点击后结果校验不通过，重新点击！' % running_time)
                            # 为了校验执行后重新点击target，需启动截图
                            need_screenShot_tap_info = tap_info.copy()
                            need_screenShot_tap_info.update({'need_screenShot': True})
                            find_result = self.find_target_imgV4(target_name, need_screenShot_tap_info)
                            # 判断重新点击target后的结果
                            if find_result[0]:
                                # 点击target成功后将会重新执行校验
                                # print('重新点击 %s 成功' % target_name)
                                logging.info('重新点击 %s 成功' % target_name)
                                if after_tap_wait_time != 0:
                                    do_after_tap_wait_time = int(after_tap_wait_time * 2)
                                    for s in range(do_after_tap_wait_time + 1):
                                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                                        half_second = s
                                        s_10, d = divmod(half_second, 10)
                                        print("\r%s %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                                        if s == do_after_tap_wait_time:
                                            pass
                                        else:
                                            time.sleep(0.5)
                                    print('\n', end='')
                                else:
                                    pass
                            # 点击target失败也将重新执行校验
                            else:
                                # print('当前界面未能找到 %s ，需重新校验 %s' % (target_name, tap_result_check_name))
                                logging.error(
                                    '当前界面未能找到 %s ，需重新校验 %s' % (target_name, tap_result_check_name))
                    if is_not_Exist_result is False:
                        break
                    else:
                        pass
                # 无校验的直接结束点击处理
                else:
                    break
            else:
                print(error_content)
                # 暂停后手动操作
                result_action = win32api.MessageBox(0, "自动执行异常请检查！请手动调整！是再次重试/否跳过此步骤:", "提醒",
                                                    win32con.MB_TOPMOST | win32con.MB_YESNO)
                if result_action == win32con.IDNO:
                    logging.info('选择了跳过 %s 步骤' % target_name)
                    break
                elif result_action == win32con.IDYES:
                    tap_info.update({'need_screenShot': True})
                    logging.info('重置截图触发为启动截图')
        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
        print('%s +++++ END +++++ [%s]\n' % (running_time, end_content))

    def is_all_ExistV2(self, target_name: str):
        """
        【判断多个指定图片是否包含在截图中】
        :param target_name:
        :return: 图片查询结果和包含详情
        """
        target_yaml = target_yaml_data[target_name]
        target_img_name = target_yaml['target_img_name_list'][0]
        rgb = target_yaml['rgb']
        threshold = target_yaml['threshold']
        screenshot_img = aircv.imread(self.screenshot_img_file)
        target_img_name_file = self.ld.target_path + target_img_name
        target_img = aircv.imread(target_img_name_file)
        result = aircv.find_all_template(screenshot_img, target_img, rgb=rgb, threshold=threshold)
        if result is None:
            return False, result
        else:
            return True, result

    def List_selectV2(self, target_name: str):
        """
        【从人形列表选择人形】
        :param self:
        :param target_name:
        """
        list_result = self.is_all_ExistV2(target_name)[1]
        # print(list_result)
        # print(len(list_result))
        if list_result is not False:
            for result in list_result:
                coordinate = result['result']
                x = coordinate[0] - 100
                y = coordinate[1] + 200
                print('选中人形 X: %s  Y: %s' % (x, y))
                self.ld.inputTap(Console_Action.index, x, y)
                time.sleep(0.5)

def tap_dict(data_dict: dict):
    """

    :param data_dict:
    """
    # print("dict:")
    for i in data_dict:
        if 'get' in i:
            get_def = data_dict[i]
            def_info_dict(i)
            drive_yaml(get_def)
            continue

def screenShot_dict(data_dict: dict):
    """

    :param data_dict:
    """
    isExist_target_key_result = False
    run_yaml = ''
    for key in data_dict:
        isExist_target_key = key
        isExist_target_value = data_dict[key]
        # print(isExist_target_key)
        if isExist_target_key == 'else':
            running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
            # print("%s 截图查询[ %s ]元素不存在！" % (running_time, data_dict.keys()))
            dict_keys = list(data_dict.keys())
            logging.warning("%s 截图查询 %s 元素不存在！" % (running_time, dict_keys))
            if isExist_target_value == 'exit':
                logging.debug('执行exit方法')
                return True
                # sys.exit()
            elif isExist_target_value == 'break':
                logging.debug('执行break方法')
                break
            else:
                running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                # print("%s 执行else方法" % running_time)
                logging.debug("执行else方法")
                tap_list(isExist_target_value)
        elif Console_Action.isExistV2(isExist_target_key)[0]:
            running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
            # print("%s 包含: %s" % (running_time, isExist_target_key))
            logging.debug("%s 包含: %s" % (running_time, isExist_target_key))
            # print(isExist_target_key, isExist_target_value)
            isExist_target_key_result = True
            run_yaml = isExist_target_value
            break
    if isExist_target_key_result:
        # 截图对象查询到后，判断后续处理对象还是
        if type(run_yaml) == dict:
            screenShot_dict(run_yaml)
        else:
            if drive_yaml(run_yaml):
                return True

def tap_info_dict(target: str, data_dict: dict):
    """

    :param target:
    :param data_dict:
    :return:
    """
    # 获取target中的默认的beginning_content与end_content值
    target_yaml = target_yaml_data[target]
    # 默认初始化的tap_info
    default_tap_info: dict = battle_yaml_data['default_tap_info']
    tap_info = default_tap_info.copy()
    tap_info.update(target_yaml)
    if data_dict is None:
        return tap_info
    else:
        tap_info.update(data_dict)
        return tap_info

def def_info_dict(data_dict: dict):
    """

    :param data_dict:
    :return:
    """
    fun: bool = True
    adv_retire: bool = True
    def_info: dict = {}
    def_info.update(data_dict)
    return def_info

def tap_list(data_target_list: list, data_info: dict = None):
    """

    :param data_info:
    :param data_target_list:
    """
    # print("list:")
    for data_target in data_target_list:
        for target in data_target:
            # 是否截图判断target存在
            if target == 'screenShot':
                if data_target[target] is None:
                    Console.ScreenShot_Ld(Console_Action.index)
                    break
                else:
                    Console.ScreenShot_Ld(Console_Action.index)
                    screenShot_dict_info = data_target[target]
                    # print('%s: %s' % (key, isExist_target[key]))
                    if screenShot_dict(screenShot_dict_info):
                        return True
            # 是否是驱动方法
            elif 'get_' in target:
                get_def_info = data_target[target]
                def_info = def_info_dict(get_def_info)
                get_def_yaml = battle_yaml_data[target]
                drive_yaml_result = drive_yaml(get_def_yaml, def_info)
                if drive_yaml_result:
                    return True
                elif drive_yaml_result is False:
                    return False
            elif 'round_' in target:
                round_times = data_target[target]
                round_def_yaml = battle_yaml_data[target]
                for i in range(round_times):
                    round_result = drive_yaml(round_def_yaml)
                    if round_result:
                        return True
                    elif round_result is False:
                        return False
                    else:
                        pass
            elif target == 'adv_retire':
                if data_info['adv_retire']:
                    tap_list(data_target[target])
            elif target == 'sleep':
                sleep_time = data_target[target]
                if sleep_time != 0:
                    do_sleep_time = int(sleep_time * 2)
                    for s in range(do_sleep_time + 1):
                        running_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
                        half_second = s
                        s_10, d = divmod(half_second, 10)
                        print("\r%s 请稍等 %-30s" % (running_time, 'o' * s_10 + '.' * d), end="")
                        if s == do_sleep_time:
                            pass
                        else:
                            time.sleep(0.5)
                    print('\n', end='')
                else:
                    pass
            elif target == 'swipe':
                swipe_info = data_target[target]
                x0 = swipe_info[0]
                y0 = swipe_info[1]
                x1 = swipe_info[2]
                y1 = swipe_info[3]
                ms = swipe_info[4]
                Console.inputSwipe(Console_Action.index, x0, y0, x1, y1, ms)
            elif target == 'loop':
                pass
            elif target == 'goto':
                pass
            else:
                if data_target[target] == 'list':
                    Console_Action.List_selectV2(target)
                else:
                    tap_info = tap_info_dict(target, data_target[target])
                    # if Ld.LdactionTapV4(target, tap_info):
                    #     return True
                    Console_Action.Action_Tap_V4(target, tap_info)

def drive_yaml(data, data_info: dict = {}, fun_return: bool = False):
    """

    :param fun_return:
    :param data_info:
    :param data:
    """
    if type(data) == dict:
        tap_dict(data)
    elif type(data) == list:
        if tap_list(data, data_info):
            return True
        if 'fun_return' in data_info.keys():
            data_info_result = data_info['fun_return']
            if data_info_result:
                return True
            else:
                return False
    else:
        pass

def test():
    """
    【调试测试类】
    :return:
    """
    # get_into_mission = battle_yaml_data['get_into_mission']
    # get_13_4 = battle_yaml_data['get_13_4']
    # drive_yaml(get_into_mission)
    # for i in range(2):
    #     drive_yaml(get_13_4)
    if Console_Action.isExistV2('end_fighting'):
        print("True")
    sys.exit()

class Job(threading.Thread):

    def __init__(self, *args, **kwargs):
        super(Job, self).__init__(*args, **kwargs)
        self.__flag = threading.Event()  # 用于暂停线程的标识
        self.__flag.set()  # 设置为True
        self.__running = threading.Event()  # 用于停止线程的标识
        self.__running.set()  # 将running设置为True

    # def run(self):
    #     while self.__running.isSet():
    #         print("thread is run")
    #         self.__flag.wait()  # 为True时立即返回, 为False时阻塞直到self.__flag为True后返回
    #         print("run = ", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    #         time.sleep(1)

    def pause(self):
        logging.debug("thread is pause")
        self.__flag.clear()  # 设置为False, 让线程阻塞

    def resume(self):
        logging.debug("thread is resume")
        self.__flag.set()  # 设置为True, 让线程停止阻塞

    def stop(self):
        logging.debug("thread is stop")
        self.__flag.set()  # 将线程从暂停状态恢复, 如何已经暂停的话
        self.__running.clear()  # 设置为False

def run_config():
    global Console, Console_Action, target_yaml_data, battle_yaml_data, get_into_mission, select_battle_name, battle_name_1, battle_name_2, \
        end_combat_1, end_combat_2
    Dc = Dnconsole(r'C:\leidian\LDPlayer9')
    Ld = Ldaction(0, Dc, 'screenshot_tmp.png')
    # print(Dc.workspace_path)
    logging.debug(Dc.workspace_path)
    target_yaml_data = Dc.YAML('target.yaml')
    # 载入关卡配置
    logging.info('开始加载关卡配置')
    battle_name = '13_4.yaml'
    battle_yaml_data = Dc.YAML(battle_name)
    get_into_mission = battle_yaml_data['get_into_mission']
    select_13_4 = battle_yaml_data['get_select_13_4']
    battle_13_4_1 = battle_yaml_data['battle_13_4_1']
    battle_13_4_2 = battle_yaml_data['battle_13_4_2']
    end_combat_1 = battle_yaml_data['end_combat_1']
    end_combat_2 = battle_yaml_data['end_combat_2']

def battle():
    global running_script, start_time
    user_action = True
    while user_action:
        try:
            runtimes = int(input("跑几圈:"))
        except ValueError:
            print("请正确输入圈数,未输入将默认执行1次")
            runtimes = 1
        try:
            runtime = float(input("跑多久(分钟）:"))
        except ValueError:
            print("请正确输入执行时间，未输入将默认执行999分钟限制")
            runtime = 999
        finally:
            runtime_min = runtime * 60
        running_script = '共 %d 次，开始执行' % runtimes
        start_time = time.time()
        if drive_yaml(get_into_mission):
            continue
        ran_time: float = 0
        if runtimes == 1:
            drive_yaml(select_battle_name)
            drive_yaml(battle_name_1)
            drive_yaml(end_combat_1)
        elif runtimes > 1:
            # 执行次数大于1，默认执行再次战斗流程
            fight_again: bool = True
            drive_yaml(select_battle_name)
            for is_runtimes in range(runtimes):
                is_runtimes_num = is_runtimes + 1
                running_script = '共 %d 次，正在执行第 %d 次' % (runtimes, is_runtimes_num)
                if is_runtimes_num > 1 and fight_again is True:
                    drive_yaml(battle_name_2)
                elif is_runtimes_num == 1 and fight_again is True:
                    drive_yaml(battle_name_1)
                elif is_runtimes_num > 1 and fight_again is False:
                    drive_yaml(battle_name_1)
                end_time = time.time()
                ran_time = end_time - start_time
                if ran_time < runtime_min and is_runtimes_num < runtimes:
                    fight_again = True
                    end_result = drive_yaml(end_combat_2)
                    # 【结束流程】校验执行了特殊方法后判断，是否是使用再次战斗执行的战斗流程
                    if end_result:
                        fight_again = False
                    end_time = time.time()
                    ran_time = end_time - start_time
                    ran_time_m, ran_time_s = divmod(ran_time, 60)
                    ran_time_h, ran_time_m = divmod(ran_time_m, 60)
                    logging.info('共 %d 次，已执行 %d 次, 已运行 %02d 时 %02d 分 %02d 秒\n' % (
                        runtimes, is_runtimes_num, ran_time_h, ran_time_m, ran_time_s))
                    # print('共 %d 次，已执行 %d 次, 已运行 %02d 时 %02d 分 %02d 秒\n' % (
                    #     runtimes, is_runtimes_num, ran_time_h, ran_time_m, ran_time_s))
                elif ran_time < runtime_min and is_runtimes_num == runtimes:
                    drive_yaml(end_combat_1)
                    end_time = time.time()
                    ran_time = end_time - start_time
                    ran_time_m, ran_time_s = divmod(ran_time, 60)
                    ran_time_h, ran_time_m = divmod(ran_time_m, 60)
                    logging.info('共 %d 次，已执行 %d 次, 已运行 %02d 时 %02d 分 %02d 秒\n' % (
                        runtimes, is_runtimes_num, ran_time_h, ran_time_m, ran_time_s))
                else:
                    drive_yaml(end_combat_1)
                    end_time = time.time()
                    ran_time = end_time - start_time
                    ran_time_m, ran_time_s = divmod(ran_time, 60)
                    ran_time_h, ran_time_m = divmod(ran_time_m, 60)
                    logging.info('共 %d 次，已执行 %d 次, 已运行 %02d 时 %02d 分 %02d 秒\n' % (
                        runtimes, is_runtimes_num, ran_time_h, ran_time_m, ran_time_s))
                    break
        result_action = win32api.MessageBox(0, "是否继续跑步机?", "提醒", win32con.MB_TOPMOST | win32con.MB_YESNO)
        # action = input('是否继续跑步机?enter后继续/输入exit退出: ')
        # if action == 'exit':
        #     break
        # else:
        #     pass
        if result_action == win32con.IDYES:
            logging.info('继续执行跑步机')
            pass
        else:
            logging.info('结束跑步机')
            break
        # os.system('pause')  # 按任意键继续

if __name__ == '__main__':
    CON_LOG = 'logger.conf'  # 配置文件路径
    logging.config.fileConfig(CON_LOG)  # '读取日志配置文件'
    logger = logging.getLogger('exampleLogger')  # 创建一个日志器logger
    workspace = os.getcwd()
    global Console, Console_Action, target_yaml_data, battle_yaml_data, get_into_mission, select_battle_name, battle_name_1, battle_name_2, \
        end_combat_1, end_combat_2, running_script, start_time
    # base_run = threading.Thread(target=run_config)
    # battle_run = threading.Thread(target=battle)
    # 基础数据和环境初始化
    run_config()
    # 启动战斗运行
    battle_run = Job(target=battle)
    battle_run.start()
    battle_run.join()

    # battle_run.pause()

    # AppPackage = 'com.sunborn.girlsfrontline.cn'
    # App = '少女前线'
    # if Dc.isrunning(0):
    #     print('模拟器已启动')
    #     time.sleep(5)
    # else:
    #     Dc.launch(0)
    #     print('模拟器启动中')
    #     time.sleep(15)
    #
    # if Dc.isruningAPP(0, AppPackage):
    #     print('%s is running' % AppPackage)
    # else:
    #     Dc.runApp(0, AppPackage)
    #     print('%s is ran' % AppPackage)
    #
    # time.sleep(10)
    # print('%s 已启动' % App)

    # # 登录（暂不支持需要输入密码）
    # Ld.LdactionTap(0, Dc, r'screenshout_tmp.png', r'sign_in.png', sleeptime=3)
    # # 开始游戏，进入主界面
    # Ld.LdactionTap(0, Dc, r'screenshout_tmp.png', r'start_game.png')

    # 调试方法
    # start_time = time.time()
    # running_script = '1'
    # test()


import tensorflow as tf
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def dd(img):
    img_array = np.array(img).astype(np.float32)
    I=img_array
    zmax=np.max(I)
    zmin=np.min(I)
    tk=(zmax+zmin)
    b=1
    m,n=I.shape
    while b==0:
        ifg=0
        ibg=0
        fnum=0
        bnum=0
        for i in range(1,m):
             for j in range(1,n):
                tmp=I(i,j)
                if tmp>=tk:
                    ifg=ifg+1
                    fnum=fnum+int(tmp)
                else:
                    ibg=ibg+1
                    bnum=bnum+int(tmp)
        zo=int(fnum/ifg)
        zb=int(bnum/ibg)
        if tk==int((zo+zb)/2):
            b=0
        else:
            tk=int((zo+zb)/2)
    return tk

img = cv2.imread("Ainstan.jpg")
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
gray = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)
img = cv2.resize(gray,(300,300))
yvzhi=dd(img)
ret1, th1 = cv2.threshold(img, yvzhi, 255, cv2.THRESH_BINARY)
print(ret1)
plt.imshow(th1,cmap=cm.gray)
plt.show()


import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import numpy as np
import random
import sys
import scipy.spatial.distance
import math
import random
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import os
import open3d as o3d
import time
from torch_geometric.nn import knn_interpolate, fps
import json
from collections import Counter, OrderedDict
from torch.optim import AdamW
from loss import *
from utils import *
from model import *

batch_size = 2
batch_size_test = 2

def train_all(model,disc,train_loader,valid_loader, n_cls,count=None, shape = None, epochs=500, smoothing= False, save=True):
    best = 0
    patience = 100
    p = 0
    if smoothing:
        sig = 0.0
    else:
        sig = 0.0
    log_var_a = torch.zeros((1,), requires_grad=True)
    log_var_b = torch.zeros((1,), requires_grad=True)

    params = ([p for p in model.parameters()] + [log_var_a] + [log_var_b])
    optimizer = torch.optim.Adam(params, lr=lr)
    optimizerD = torch.optim.Adam(disc.parameters(), lr=lr2)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, step*len(train_loader), eta_min=1e-4)
    schedulerD = torch.optim.lr_scheduler.CosineAnnealingLR(optimizerD, step*len(train_loader), eta_min=1e-4)

    for epoch in range(epochs):
        e = epoch
        model.train()
        running_loss = 0
        for i, data in enumerate(train_loader):
            inputs, labels = data['points'].to(device), data['labels'].to(device)
            #UPDATE DISCRIMINATOR
            optimizerD.zero_grad()

            labels = labels.view(batch_size, -1)

            labels2 = one_hot(labels,n_cls).to(device)
            noise = torch.tensor(np.random.normal(0, sig, labels2.shape)).to(device)
            labels2 = labels2+noise

            inputs2 = torch.cat((inputs, labels2.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
            inputs2 = inputs2.transpose(1,2)
            real_image = inputs2
            output,em1 = disc(inputs2.float())
            ls_real = disc_loss(output,num=1, smoothing = smoothing)
            ls_real.backward(retain_graph=True)

            #Fake image
            with torch.no_grad():
                output = model(inputs.transpose(1,2))
            output_pred = output
            _,pred = torch.max(output.data,1)
            pred = pred.view(batch_size,-1)
            pred = one_hot(pred,n_cls)
            noise = torch.tensor(np.random.normal(0, sig, pred.shape))#.to(device)
            pred += noise
            pred = pred.to(device)
            inputs2 = torch.cat((inputs, pred.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
            inputs2 = inputs2.transpose(1,2)
            output,em2 = disc(inputs2.float())
            ls_fake = disc_loss(output, num=0, smoothing = smoothing)
            ls_fake.backward()

            ls_D = (ls_real + ls_fake)/2
            optimizerD.step()
            schedulerD.step()

            #UPDATE GENERATOR
            #Real image
            optimizer.zero_grad()
            #real
            inputs2 = torch.cat((inputs, labels2.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
            inputs2 = inputs2.transpose(1,2)
            real_image = inputs2
            #fake
            output = model(inputs.transpose(1,2))
            output_pred = output
            output2 = output.clone()
            _,pred = torch.max(output.data,1)
            pred = pred.view(batch_size,-1)
            pred = one_hot(pred,n_cls)
            noise = torch.tensor(np.random.normal(0, sig, pred.shape))#.to(device)
            pred += noise
            pred = pred.to(device)
            inputs2 = torch.cat((inputs, pred.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
            inputs2 = inputs2.transpose(1,2)

            _,em1 = disc(real_image.float())
            _,em2 = disc(inputs2.float())

            output_pred = [output_pred,em2]
            labels = [labels,em1]
            total_ls, emb_ls = multitask_loss(output_pred,labels,n_cls,[log_var_a,log_var_b])
            total_ls.backward()

            #if i%30 == 0:
            #    print("Embedding Loss: {}".format(emb_ls))
            optimizer.step()
            scheduler.step()

        std_1 = torch.exp(-log_var_a)**0.5
        std_2 = torch.exp(-log_var_b)**0.5

        model.eval()
        total = 0
        correct =0
        total2 = 0
        correct2 =0
        loss =[]
        loss2 = []
        valid_iou = []
        with torch.no_grad():
            for i, data in enumerate(valid_loader):
                inputs, labels = data['points'].to(device), data['labels'].to(device)

                #Discriminator
                #Real
                labels = labels.view(batch_size, -1)

                labels2 = one_hot(labels,n_cls).to(device)
                noise = torch.tensor(np.random.normal(0, sig, labels2.shape)).to(device)
                labels2 = labels2+noise

                inputs2 = torch.cat((inputs, labels2.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
                inputs2 = inputs2.transpose(1,2)
                real_image = inputs2
                output,em1 = disc(inputs2.float())
                _,pred2 = torch.max(output.data,1)
                ls_real = disc_loss(output,num=1)
                total2 += batch_size_test
                correct2 += (pred2 == 1).sum().item()

                #Fake
                with torch.no_grad():
                    outputs = model(inputs.transpose(1,2))
                output_pred = outputs
                _,pred = torch.max(outputs.data,1)
                pred = pred.view(batch_size,-1)
                pred = one_hot(pred,n_cls)
                noise = torch.tensor(np.random.normal(0, sig, pred.shape))#.to(device)
                pred += noise
                pred = pred.to(device)
                inputs2 = torch.cat((inputs, pred.type(torch.float64).reshape(batch_size,2048,n_cls)),axis=2)
                inputs2 = inputs2.transpose(1,2)
                output,em2 = disc(inputs2.float())
                _,pred2 = torch.max(output.data,1)
                ls_fake = disc_loss(output, num=0)

                ls_total = (ls_real + ls_fake)/2
                loss2.append(ls_total.item())
                total2 += batch_size_test
                correct2 += (pred2 == 0).sum().item()

                #Generator
                outputs = model(inputs.transpose(1,2))
                ls = gagcn_loss(outputs,labels,n_cls, smoothing = True)

                loss.append(ls.item())
                _,pred = torch.max(outputs.data,1) #why? because the first return is for the values(probabilities)
                total += labels.size(0) * labels.size(1)
                correct += (pred == labels).sum().item()
                pred = pred.cpu().detach().numpy()
                ious = calculate_shape_IoU(pred,labels.cpu().detach().numpy(),n_cls)
                valid_iou.append(np.mean(ious))

            #Generaotr
            acc = 100. * correct / total
            loss_final = np.mean(loss)
            iou_final = np.mean(valid_iou)
            print("Epoch: {}".format(e))
            print("Generator: The validation accuracy: {:.6f} and the validation loss: {:.6f} and the validation iou: {:.6f}".format(acc,loss_final,iou_final))
            print("std1 {} std2 {}".format(std_1,std_2))
            #Discriminator
            acc = 100. * correct2 / total2
            loss_final = np.mean(loss2)
            print("Discriminator: The validation accuracy: {:.6f} and the validation loss: {:.6f}".format(acc,loss_final))

            if best < iou_final:
                best = iou_final
                p = 0
                if save:
                    torch.save(model.state_dict(), "/home/edshkim98/shkim/pretrained_group_final/partseg_"+str(shape)+'_'+str(count)+".pth")
                else:
                    pass
            else:
                p +=1
                print("Patience: {} / {}".format(p,patience))
                if p == patience:
                    print("Stopped due to convergence")
                    return best
            print("Valid best iou: ",best)
    return best

if __name__ == "__main__":
    dataset = {"02691156": "Airplane", "02773838": "Bag", "02954340": "Cap", "02958343": "Car", "03001627": "Chair",
           "03261776": "Earphone", "03467517": "Guitar", "03624134": "Knife", "03636649": "Lamp", "03642806": "Laptop",
           "03790512": "Motorbike", "03797390": "Mug", "03948459": "Pistol", "04099429": "Rocket", "04225987": "Skateboard",
           "04379243": "Table"}
    seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3]
    print("Training on ShapeNet dataset")
    print("#############################")
    total_final = 0
    total_num = 0
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    for o in range(15,len(dataset.keys())):
        shape = list(dataset.keys())[-3]
        path = r'/home/edshkim98/shkim/shape_data/'+shape+'/'#+shape
        new_path =path+'/points_label'
        n_cls = seg_num[-3]

        valid_files = '/home/edshkim98/shkim/shape_data/train_test_split/shuffled_val_file_list.json'
        train_files = '/home/edshkim98/shkim/shape_data/train_test_split/shuffled_train_file_list.json'
        test_files = '/home/edshkim98/shkim/shape_data/train_test_split/shuffled_test_file_list.json'

        torch.manual_seed(42)
        torch.cuda.manual_seed(42)
        random.seed(42)
        np.random.seed(42)

        train_files = read_json(train_files)
        val_files = read_json(valid_files)
        test_files = read_json(test_files)

        train_xs = []
        val = []
        test = []
        for i in train_files:
            if shape in i:
                train_xs.append(i)
        for i in val_files:
            if shape in i:
                val.append(i)
        for i in test_files:
            if shape in i:
                test.append(i)
        total = train_xs+val+test
        np.random.shuffle(total)
        #Train
        train_xs = train_xs+val#np.array(total)[idx]
        #Val&Test
        val = test
        train_dataset = CustomDataset(path,train_xs, transform =train_transforms,valid = False)
        val_dataset = CustomDataset(path, val,transform=None, valid=True)

        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last = True, num_workers =2)
        val_loader = DataLoader(val_dataset, batch_size=batch_size_test, shuffle=False, drop_last = True, num_workers =2)

        print('######### Dataset class created #########')
        print("Shape: {} Batch size: {}".format(shape,batch_size))
        print('Number of images: ', len(train_xs)+len(val))
        print('Sample image shape: ', train_dataset[0]['points'].shape)
        print('train size in no of batch: ',len(train_loader))
        print("test size in no of batch: ",len(val_loader))
        print('train size: ',len(train_loader)*batch_size)
        print("valid size: ",len(val_loader)*batch_size_test)

            #########################

            ###########################
        best = 0
        total = len(val_loader)*batch_size_test

        if len(train_loader)*batch_size > 1500:
            cnts = 1
        else:
            cnts=2

        for cnt in range(cnts):
            smoothing=False
            bias = True
            dilation_rate = [1,3,6,6,6,1]
            pyramid = [False,False,False,False,False]
            dilation = 'order'
            lr = 0.005
            lr2 = 0.01
            step = 5

            se = False
            gagcn = GAGCN(n_cls,k=16,pyramid=pyramid,dilation='order', dilation_rate = dilation_rate, norm='group', bias= bias)
            torch.cuda.set_device(0)
            gagcn.to(device)
            pytorch_total_params = sum(p.numel() for p in gagcn.parameters())
            print("Number of parameters (segmentation): ", pytorch_total_params)

            se = False
            disc = Discriminator(k=12, dilation='order', dilation_rate= [1,3,3,3], n_cls= n_cls, out_dims= [128,256,256,512],pyramid=[False,False,False,False],norm='group')
            disc.to(device)
            pytorch_total_params = sum(p.numel() for p in disc.parameters())
            print("Number of parameters (discriminator): ", pytorch_total_params)
            print("Disc smoothing: ",smoothing)

            best_curr = train_all(gagcn,disc,train_loader, val_loader, n_cls, count=cnt,shape=shape, smoothing = smoothing,save=False)
            if best_curr > best:
                best = best_curr
                print("###############################################")
                print("Best Setting-> Shape {} IoU {} count {}".format(shape, best, cnt))
                print("Best: {}%".format(best*100))
                print("###############################################")