# -*- coding: utf-8 -*-
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import dataclasses
import json  # type: ignore
import re
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
import warnings

from google.api_core import (
    gapic_v1,
    operations_v1,
    path_template,
    rest_helpers,
    rest_streaming,
)
from google.api_core import exceptions as core_exceptions
from google.api_core import retry as retries
from google.auth import credentials as ga_credentials  # type: ignore
from google.auth.transport.grpc import SslCredentials  # type: ignore
from google.auth.transport.requests import AuthorizedSession  # type: ignore
from google.protobuf import json_format
import grpc  # type: ignore
from requests import __version__ as requests_version

try:
    OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault]
except AttributeError:  # pragma: NO COVER
    OptionalRetry = Union[retries.Retry, object]  # type: ignore

from google.longrunning import operations_pb2  # type: ignore

from google.analytics.data_v1alpha.types import analytics_data_api

from .base import AlphaAnalyticsDataTransport
from .base import DEFAULT_CLIENT_INFO as BASE_DEFAULT_CLIENT_INFO

DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
    gapic_version=BASE_DEFAULT_CLIENT_INFO.gapic_version,
    grpc_version=None,
    rest_version=requests_version,
)

class AlphaAnalyticsDataRestInterceptor:
    """Interceptor for AlphaAnalyticsData.

    Interceptors are used to manipulate requests, request metadata, and responses
    in arbitrary ways.
    Example use cases include:
    * Logging
    * Verifying requests according to service or custom semantics
    * Stripping extraneous information from responses

    These use cases and more can be enabled by injecting an
    instance of a custom subclass when constructing the AlphaAnalyticsDataRestTransport.

    .. code-block:: python
        class MyCustomAlphaAnalyticsDataInterceptor(AlphaAnalyticsDataRestInterceptor):
            def pre_create_audience_list(self, request, metadata):
                logging.log(f"Received request: {request}")
                return request, metadata

            def post_create_audience_list(self, response):
                logging.log(f"Received response: {response}")
                return response

            def pre_get_audience_list(self, request, metadata):
                logging.log(f"Received request: {request}")
                return request, metadata

            def post_get_audience_list(self, response):
                logging.log(f"Received response: {response}")
                return response

            def pre_list_audience_lists(self, request, metadata):
                logging.log(f"Received request: {request}")
                return request, metadata

            def post_list_audience_lists(self, response):
                logging.log(f"Received response: {response}")
                return response

            def pre_query_audience_list(self, request, metadata):
                logging.log(f"Received request: {request}")
                return request, metadata

            def post_query_audience_list(self, response):
                logging.log(f"Received response: {response}")
                return response

            def pre_run_funnel_report(self, request, metadata):
                logging.log(f"Received request: {request}")
                return request, metadata

            def post_run_funnel_report(self, response):
                logging.log(f"Received response: {response}")
                return response

        transport = AlphaAnalyticsDataRestTransport(interceptor=MyCustomAlphaAnalyticsDataInterceptor())
        client = AlphaAnalyticsDataClient(transport=transport)

    """

    def pre_create_audience_list(
        self,
        request: analytics_data_api.CreateAudienceListRequest,
        metadata: Sequence[Tuple[str, str]],
    ) -> Tuple[analytics_data_api.CreateAudienceListRequest, Sequence[Tuple[str, str]]]:
        """Pre-rpc interceptor for create_audience_list

        Override in a subclass to manipulate the request or metadata
        before they are sent to the AlphaAnalyticsData server.
        """
        return request, metadata

    def post_create_audience_list(
        self, response: operations_pb2.Operation
    ) -> operations_pb2.Operation:
        """Post-rpc interceptor for create_audience_list

        Override in a subclass to manipulate the response
        after it is returned by the AlphaAnalyticsData server but before
        it is returned to user code.
        """
        return response

    def pre_get_audience_list(
        self,
        request: analytics_data_api.GetAudienceListRequest,
        metadata: Sequence[Tuple[str, str]],
    ) -> Tuple[analytics_data_api.GetAudienceListRequest, Sequence[Tuple[str, str]]]:
        """Pre-rpc interceptor for get_audience_list

        Override in a subclass to manipulate the request or metadata
        before they are sent to the AlphaAnalyticsData server.
        """
        return request, metadata

    def post_get_audience_list(
        self, response: analytics_data_api.AudienceList
    ) -> analytics_data_api.AudienceList:
        """Post-rpc interceptor for get_audience_list

        Override in a subclass to manipulate the response
        after it is returned by the AlphaAnalyticsData server but before
        it is returned to user code.
        """
        return response

    def pre_list_audience_lists(
        self,
        request: analytics_data_api.ListAudienceListsRequest,
        metadata: Sequence[Tuple[str, str]],
    ) -> Tuple[analytics_data_api.ListAudienceListsRequest, Sequence[Tuple[str, str]]]:
        """Pre-rpc interceptor for list_audience_lists

        Override in a subclass to manipulate the request or metadata
        before they are sent to the AlphaAnalyticsData server.
        """
        return request, metadata

    def post_list_audience_lists(
        self, response: analytics_data_api.ListAudienceListsResponse
    ) -> analytics_data_api.ListAudienceListsResponse:
        """Post-rpc interceptor for list_audience_lists

        Override in a subclass to manipulate the response
        after it is returned by the AlphaAnalyticsData server but before
        it is returned to user code.
        """
        return response

    def pre_query_audience_list(
        self,
        request: analytics_data_api.QueryAudienceListRequest,
        metadata: Sequence[Tuple[str, str]],
    ) -> Tuple[analytics_data_api.QueryAudienceListRequest, Sequence[Tuple[str, str]]]:
        """Pre-rpc interceptor for query_audience_list

        Override in a subclass to manipulate the request or metadata
        before they are sent to the AlphaAnalyticsData server.
        """
        return request, metadata

    def post_query_audience_list(
        self, response: analytics_data_api.QueryAudienceListResponse
    ) -> analytics_data_api.QueryAudienceListResponse:
        """Post-rpc interceptor for query_audience_list

        Override in a subclass to manipulate the response
        after it is returned by the AlphaAnalyticsData server but before
        it is returned to user code.
        """
        return response

    def pre_run_funnel_report(
        self,
        request: analytics_data_api.RunFunnelReportRequest,
        metadata: Sequence[Tuple[str, str]],
    ) -> Tuple[analytics_data_api.RunFunnelReportRequest, Sequence[Tuple[str, str]]]:
        """Pre-rpc interceptor for run_funnel_report

        Override in a subclass to manipulate the request or metadata
        before they are sent to the AlphaAnalyticsData server.
        """
        return request, metadata

    def post_run_funnel_report(
        self, response: analytics_data_api.RunFunnelReportResponse
    ) -> analytics_data_api.RunFunnelReportResponse:
        """Post-rpc interceptor for run_funnel_report

        Override in a subclass to manipulate the response
        after it is returned by the AlphaAnalyticsData server but before
        it is returned to user code.
        """
        return response

@dataclasses.dataclass
class AlphaAnalyticsDataRestStub:
    _session: AuthorizedSession
    _host: str
    _interceptor: AlphaAnalyticsDataRestInterceptor

class AlphaAnalyticsDataRestTransport(AlphaAnalyticsDataTransport):
    """REST backend transport for AlphaAnalyticsData.

    Google Analytics reporting data service.

    This class defines the same methods as the primary client, so the
    primary client can load the underlying transport implementation
    and call it.

    It sends JSON representations of protocol buffers over HTTP/1.1

    """

    def __init__(
        self,
        *,
        host: str = "analyticsdata.googleapis.com",
        credentials: Optional[ga_credentials.Credentials] = None,
        credentials_file: Optional[str] = None,
        scopes: Optional[Sequence[str]] = None,
        client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None,
        quota_project_id: Optional[str] = None,
        client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
        always_use_jwt_access: Optional[bool] = False,
        url_scheme: str = "https",
        interceptor: Optional[AlphaAnalyticsDataRestInterceptor] = None,
        api_audience: Optional[str] = None,
    ) -> None:
        """Instantiate the transport.

        Args:
            host (Optional[str]):
                 The hostname to connect to.
            credentials (Optional[google.auth.credentials.Credentials]): The
                authorization credentials to attach to requests. These
                credentials identify the application to the service; if none
                are specified, the client will attempt to ascertain the
                credentials from the environment.

            credentials_file (Optional[str]): A file with credentials that can
                be loaded with :func:`google.auth.load_credentials_from_file`.
                This argument is ignored if ``channel`` is provided.
            scopes (Optional(Sequence[str])): A list of scopes. This argument is
                ignored if ``channel`` is provided.
            client_cert_source_for_mtls (Callable[[], Tuple[bytes, bytes]]): Client
                certificate to configure mutual TLS HTTP channel. It is ignored
                if ``channel`` is provided.
            quota_project_id (Optional[str]): An optional project to use for billing
                and quota.
            client_info (google.api_core.gapic_v1.client_info.ClientInfo):
                The client info used to send a user-agent string along with
                API requests. If ``None``, then default info will be used.
                Generally, you only need to set this if you are developing
                your own client library.
            always_use_jwt_access (Optional[bool]): Whether self signed JWT should
                be used for service account credentials.
            url_scheme: the protocol scheme for the API endpoint.  Normally
                "https", but for testing or local servers,
                "http" can be specified.
        """
        # Run the base constructor
        # TODO(yon-mg): resolve other ctor params i.e. scopes, quota, etc.
        # TODO: When custom host (api_endpoint) is set, `scopes` must *also* be set on the
        # credentials object
        maybe_url_match = re.match("^(?P<scheme>http(?:s)?://)?(?P<host>.*)$", host)
        if maybe_url_match is None:
            raise ValueError(
                f"Unexpected hostname structure: {host}"
            )  # pragma: NO COVER

        url_match_items = maybe_url_match.groupdict()

        host = f"{url_scheme}://{host}" if not url_match_items["scheme"] else host

        super().__init__(
            host=host,
            credentials=credentials,
            client_info=client_info,
            always_use_jwt_access=always_use_jwt_access,
            api_audience=api_audience,
        )
        self._session = AuthorizedSession(
            self._credentials, default_host=self.DEFAULT_HOST
        )
        self._operations_client: Optional[operations_v1.AbstractOperationsClient] = None
        if client_cert_source_for_mtls:
            self._session.configure_mtls_channel(client_cert_source_for_mtls)
        self._interceptor = interceptor or AlphaAnalyticsDataRestInterceptor()
        self._prep_wrapped_messages(client_info)

    @property
    def operations_client(self) -> operations_v1.AbstractOperationsClient:
        """Create the client designed to process long-running operations.

        This property caches on the instance; repeated calls return the same
        client.
        """
        # Only create a new client if we do not already have one.
        if self._operations_client is None:
            http_options: Dict[str, List[Dict[str, str]]] = {}

            rest_transport = operations_v1.OperationsRestTransport(
                host=self._host,
                # use the credentials which are saved
                credentials=self._credentials,
                scopes=self._scopes,
                http_options=http_options,
                path_prefix="v1alpha",
            )

            self._operations_client = operations_v1.AbstractOperationsClient(
                transport=rest_transport
            )

        # Return the client from cache.
        return self._operations_client

    class _CreateAudienceList(AlphaAnalyticsDataRestStub):
        def __hash__(self):
            return hash("CreateAudienceList")

        __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {}

        @classmethod
        def _get_unset_required_fields(cls, message_dict):
            return {
                k: v
                for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items()
                if k not in message_dict
            }

        def __call__(
            self,
            request: analytics_data_api.CreateAudienceListRequest,
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
        ) -> operations_pb2.Operation:
            r"""Call the create audience list method over HTTP.

            Args:
                request (~.analytics_data_api.CreateAudienceListRequest):
                    The request object. A request to create a new audience
                list.
                retry (google.api_core.retry.Retry): Designation of what errors, if any,
                    should be retried.
                timeout (float): The timeout for this request.
                metadata (Sequence[Tuple[str, str]]): Strings which should be
                    sent along with the request as metadata.

            Returns:
                ~.operations_pb2.Operation:
                    This resource represents a
                long-running operation that is the
                result of a network API call.

            """

            http_options: List[Dict[str, str]] = [
                {
                    "method": "post",
                    "uri": "/v1alpha/{parent=properties/*}/audienceLists",
                    "body": "audience_list",
                },
            ]
            request, metadata = self._interceptor.pre_create_audience_list(
                request, metadata
            )
            pb_request = analytics_data_api.CreateAudienceListRequest.pb(request)
            transcoded_request = path_template.transcode(http_options, pb_request)

            # Jsonify the request body

            body = json_format.MessageToJson(
                transcoded_request["body"],
                including_default_value_fields=False,
                use_integers_for_enums=True,
            )
            uri = transcoded_request["uri"]
            method = transcoded_request["method"]

            # Jsonify the query params
            query_params = json.loads(
                json_format.MessageToJson(
                    transcoded_request["query_params"],
                    including_default_value_fields=False,
                    use_integers_for_enums=True,
                )
            )
            query_params.update(self._get_unset_required_fields(query_params))

            query_params["$alt"] = "json;enum-encoding=int"

            # Send the request
            headers = dict(metadata)
            headers["Content-Type"] = "application/json"
            response = getattr(self._session, method)(
                "{host}{uri}".format(host=self._host, uri=uri),
                timeout=timeout,
                headers=headers,
                params=rest_helpers.flatten_query_params(query_params, strict=True),
                data=body,
            )

            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
            # subclass.
            if response.status_code >= 400:
                raise core_exceptions.from_http_response(response)

            # Return the response
            resp = operations_pb2.Operation()
            json_format.Parse(response.content, resp, ignore_unknown_fields=True)
            resp = self._interceptor.post_create_audience_list(resp)
            return resp

    class _GetAudienceList(AlphaAnalyticsDataRestStub):
        def __hash__(self):
            return hash("GetAudienceList")

        __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {}

        @classmethod
        def _get_unset_required_fields(cls, message_dict):
            return {
                k: v
                for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items()
                if k not in message_dict
            }

        def __call__(
            self,
            request: analytics_data_api.GetAudienceListRequest,
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
        ) -> analytics_data_api.AudienceList:
            r"""Call the get audience list method over HTTP.

            Args:
                request (~.analytics_data_api.GetAudienceListRequest):
                    The request object. A request to retrieve configuration
                metadata about a specific audience list.
                retry (google.api_core.retry.Retry): Designation of what errors, if any,
                    should be retried.
                timeout (float): The timeout for this request.
                metadata (Sequence[Tuple[str, str]]): Strings which should be
                    sent along with the request as metadata.

            Returns:
                ~.analytics_data_api.AudienceList:
                    An audience list is a list of users
                in an audience at the time of the list's
                creation. One audience may have multiple
                audience lists created for different
                days.

            """

            http_options: List[Dict[str, str]] = [
                {
                    "method": "get",
                    "uri": "/v1alpha/{name=properties/*/audienceLists/*}",
                },
            ]
            request, metadata = self._interceptor.pre_get_audience_list(
                request, metadata
            )
            pb_request = analytics_data_api.GetAudienceListRequest.pb(request)
            transcoded_request = path_template.transcode(http_options, pb_request)

            uri = transcoded_request["uri"]
            method = transcoded_request["method"]

            # Jsonify the query params
            query_params = json.loads(
                json_format.MessageToJson(
                    transcoded_request["query_params"],
                    including_default_value_fields=False,
                    use_integers_for_enums=True,
                )
            )
            query_params.update(self._get_unset_required_fields(query_params))

            query_params["$alt"] = "json;enum-encoding=int"

            # Send the request
            headers = dict(metadata)
            headers["Content-Type"] = "application/json"
            response = getattr(self._session, method)(
                "{host}{uri}".format(host=self._host, uri=uri),
                timeout=timeout,
                headers=headers,
                params=rest_helpers.flatten_query_params(query_params, strict=True),
            )

            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
            # subclass.
            if response.status_code >= 400:
                raise core_exceptions.from_http_response(response)

            # Return the response
            resp = analytics_data_api.AudienceList()
            pb_resp = analytics_data_api.AudienceList.pb(resp)

            json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)
            resp = self._interceptor.post_get_audience_list(resp)
            return resp

    class _ListAudienceLists(AlphaAnalyticsDataRestStub):
        def __hash__(self):
            return hash("ListAudienceLists")

        __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {}

        @classmethod
        def _get_unset_required_fields(cls, message_dict):
            return {
                k: v
                for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items()
                if k not in message_dict
            }

        def __call__(
            self,
            request: analytics_data_api.ListAudienceListsRequest,
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
        ) -> analytics_data_api.ListAudienceListsResponse:
            r"""Call the list audience lists method over HTTP.

            Args:
                request (~.analytics_data_api.ListAudienceListsRequest):
                    The request object. A request to list all audience lists
                for a property.
                retry (google.api_core.retry.Retry): Designation of what errors, if any,
                    should be retried.
                timeout (float): The timeout for this request.
                metadata (Sequence[Tuple[str, str]]): Strings which should be
                    sent along with the request as metadata.

            Returns:
                ~.analytics_data_api.ListAudienceListsResponse:
                    A list of all audience lists for a
                property.

            """

            http_options: List[Dict[str, str]] = [
                {
                    "method": "get",
                    "uri": "/v1alpha/{parent=properties/*}/audienceLists",
                },
            ]
            request, metadata = self._interceptor.pre_list_audience_lists(
                request, metadata
            )
            pb_request = analytics_data_api.ListAudienceListsRequest.pb(request)
            transcoded_request = path_template.transcode(http_options, pb_request)

            uri = transcoded_request["uri"]
            method = transcoded_request["method"]

            # Jsonify the query params
            query_params = json.loads(
                json_format.MessageToJson(
                    transcoded_request["query_params"],
                    including_default_value_fields=False,
                    use_integers_for_enums=True,
                )
            )
            query_params.update(self._get_unset_required_fields(query_params))

            query_params["$alt"] = "json;enum-encoding=int"

            # Send the request
            headers = dict(metadata)
            headers["Content-Type"] = "application/json"
            response = getattr(self._session, method)(
                "{host}{uri}".format(host=self._host, uri=uri),
                timeout=timeout,
                headers=headers,
                params=rest_helpers.flatten_query_params(query_params, strict=True),
            )

            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
            # subclass.
            if response.status_code >= 400:
                raise core_exceptions.from_http_response(response)

            # Return the response
            resp = analytics_data_api.ListAudienceListsResponse()
            pb_resp = analytics_data_api.ListAudienceListsResponse.pb(resp)

            json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)
            resp = self._interceptor.post_list_audience_lists(resp)
            return resp

    class _QueryAudienceList(AlphaAnalyticsDataRestStub):
        def __hash__(self):
            return hash("QueryAudienceList")

        def __call__(
            self,
            request: analytics_data_api.QueryAudienceListRequest,
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
        ) -> analytics_data_api.QueryAudienceListResponse:
            r"""Call the query audience list method over HTTP.

            Args:
                request (~.analytics_data_api.QueryAudienceListRequest):
                    The request object. A request to list users in an
                audience list.
                retry (google.api_core.retry.Retry): Designation of what errors, if any,
                    should be retried.
                timeout (float): The timeout for this request.
                metadata (Sequence[Tuple[str, str]]): Strings which should be
                    sent along with the request as metadata.

            Returns:
                ~.analytics_data_api.QueryAudienceListResponse:
                    A list of users in an audience list.
            """

            http_options: List[Dict[str, str]] = [
                {
                    "method": "post",
                    "uri": "/v1alpha/{name=properties/*/audienceLists/*}:query",
                    "body": "*",
                },
            ]
            request, metadata = self._interceptor.pre_query_audience_list(
                request, metadata
            )
            pb_request = analytics_data_api.QueryAudienceListRequest.pb(request)
            transcoded_request = path_template.transcode(http_options, pb_request)

            # Jsonify the request body

            body = json_format.MessageToJson(
                transcoded_request["body"],
                including_default_value_fields=False,
                use_integers_for_enums=True,
            )
            uri = transcoded_request["uri"]
            method = transcoded_request["method"]

            # Jsonify the query params
            query_params = json.loads(
                json_format.MessageToJson(
                    transcoded_request["query_params"],
                    including_default_value_fields=False,
                    use_integers_for_enums=True,
                )
            )

            query_params["$alt"] = "json;enum-encoding=int"

            # Send the request
            headers = dict(metadata)
            headers["Content-Type"] = "application/json"
            response = getattr(self._session, method)(
                "{host}{uri}".format(host=self._host, uri=uri),
                timeout=timeout,
                headers=headers,
                params=rest_helpers.flatten_query_params(query_params, strict=True),
                data=body,
            )

            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
            # subclass.
            if response.status_code >= 400:
                raise core_exceptions.from_http_response(response)

            # Return the response
            resp = analytics_data_api.QueryAudienceListResponse()
            pb_resp = analytics_data_api.QueryAudienceListResponse.pb(resp)

            json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)
            resp = self._interceptor.post_query_audience_list(resp)
            return resp

    class _RunFunnelReport(AlphaAnalyticsDataRestStub):
        def __hash__(self):
            return hash("RunFunnelReport")

        def __call__(
            self,
            request: analytics_data_api.RunFunnelReportRequest,
            *,
            retry: OptionalRetry = gapic_v1.method.DEFAULT,
            timeout: Optional[float] = None,
            metadata: Sequence[Tuple[str, str]] = (),
        ) -> analytics_data_api.RunFunnelReportResponse:
            r"""Call the run funnel report method over HTTP.

            Args:
                request (~.analytics_data_api.RunFunnelReportRequest):
                    The request object. The request for a funnel report.
                retry (google.api_core.retry.Retry): Designation of what errors, if any,
                    should be retried.
                timeout (float): The timeout for this request.
                metadata (Sequence[Tuple[str, str]]): Strings which should be
                    sent along with the request as metadata.

            Returns:
                ~.analytics_data_api.RunFunnelReportResponse:
                    The funnel report response contains
                two sub reports. The two sub reports are
                different combinations of dimensions and
                metrics.

            """

            http_options: List[Dict[str, str]] = [
                {
                    "method": "post",
                    "uri": "/v1alpha/{property=properties/*}:runFunnelReport",
                    "body": "*",
                },
            ]
            request, metadata = self._interceptor.pre_run_funnel_report(
                request, metadata
            )
            pb_request = analytics_data_api.RunFunnelReportRequest.pb(request)
            transcoded_request = path_template.transcode(http_options, pb_request)

            # Jsonify the request body

            body = json_format.MessageToJson(
                transcoded_request["body"],
                including_default_value_fields=False,
                use_integers_for_enums=True,
            )
            uri = transcoded_request["uri"]
            method = transcoded_request["method"]

            # Jsonify the query params
            query_params = json.loads(
                json_format.MessageToJson(
                    transcoded_request["query_params"],
                    including_default_value_fields=False,
                    use_integers_for_enums=True,
                )
            )

            query_params["$alt"] = "json;enum-encoding=int"

            # Send the request
            headers = dict(metadata)
            headers["Content-Type"] = "application/json"
            response = getattr(self._session, method)(
                "{host}{uri}".format(host=self._host, uri=uri),
                timeout=timeout,
                headers=headers,
                params=rest_helpers.flatten_query_params(query_params, strict=True),
                data=body,
            )

            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
            # subclass.
            if response.status_code >= 400:
                raise core_exceptions.from_http_response(response)

            # Return the response
            resp = analytics_data_api.RunFunnelReportResponse()
            pb_resp = analytics_data_api.RunFunnelReportResponse.pb(resp)

            json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True)
            resp = self._interceptor.post_run_funnel_report(resp)
            return resp

    @property
    def create_audience_list(
        self,
    ) -> Callable[
        [analytics_data_api.CreateAudienceListRequest], operations_pb2.Operation
    ]:
        # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here.
        # In C++ this would require a dynamic_cast
        return self._CreateAudienceList(self._session, self._host, self._interceptor)  # type: ignore

    @property
    def get_audience_list(
        self,
    ) -> Callable[
        [analytics_data_api.GetAudienceListRequest], analytics_data_api.AudienceList
    ]:
        # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here.
        # In C++ this would require a dynamic_cast
        return self._GetAudienceList(self._session, self._host, self._interceptor)  # type: ignore

    @property
    def list_audience_lists(
        self,
    ) -> Callable[
        [analytics_data_api.ListAudienceListsRequest],
        analytics_data_api.ListAudienceListsResponse,
    ]:
        # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here.
        # In C++ this would require a dynamic_cast
        return self._ListAudienceLists(self._session, self._host, self._interceptor)  # type: ignore

    @property
    def query_audience_list(
        self,
    ) -> Callable[
        [analytics_data_api.QueryAudienceListRequest],
        analytics_data_api.QueryAudienceListResponse,
    ]:
        # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here.
        # In C++ this would require a dynamic_cast
        return self._QueryAudienceList(self._session, self._host, self._interceptor)  # type: ignore

    @property
    def run_funnel_report(
        self,
    ) -> Callable[
        [analytics_data_api.RunFunnelReportRequest],
        analytics_data_api.RunFunnelReportResponse,
    ]:
        # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here.
        # In C++ this would require a dynamic_cast
        return self._RunFunnelReport(self._session, self._host, self._interceptor)  # type: ignore

    @property
    def kind(self) -> str:
        return "rest"

    def close(self):
        self._session.close()

__all__ = ("AlphaAnalyticsDataRestTransport",)


#pylint: disable=E1101
import os
import os.path
import numpy as np
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

import torch
import torchvision

import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader, ConcatDataset
from torch.utils.data.sampler import SubsetRandomSampler

import loaddataset as lds

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from scale_cnn.convolution import ScaleConvolution
from scale_cnn.pooling import ScalePool

from architectures import SiCNN_3

from functions import train, test, plot_figures
from rescale import RandomRescale
import pickle

device =  torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

nb_epochs=150
learning_rate = 0.0001
batch_size = 128
batch_log = 70
repeats = 6

f_in = 1
size=5
ratio=2**(2/3)
nratio=3
srange=0
padding=0

log = open("mnist_range_sr0_log.pickle", "wb")

parameters = {
    "epochs": nb_epochs,
    "learning rate": learning_rate,
    "batch size": batch_size,
    "repetitions": repeats,
    "size": size,
    "ratio": ratio,
    "nratio": nratio,
    "srange": srange
}
pickle.dump(parameters, log)

scales = [(1.0, 1.0), (0.9, 1.1), (0.8, 1.2), (0.6, 1.4), (0.5, 1.5), (0.4, 1.6), (0.3, 1.7)]

pickle.dump(scales, log)

criterion = nn.CrossEntropyLoss()

avg_test_losses = []
avg_test_accs = []
std_test_losses = []
std_test_accs = []

for scale in scales:
    uniform = transforms.Compose([
                transforms.Resize(40), RandomRescale(size = 40, scales = scale, sampling = "uniform"),
                transforms.ToTensor(), transforms.Normalize((0.1307,),(0.3081,))])

    root = './mnistdata'
    if not os.path.exists(root):
        os.mkdir(root)

    train_set = datasets.MNIST(root=root, train=True, transform=uniform, download=True)
    train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True, num_workers=1, pin_memory=True)

    test_set = datasets.MNIST(root=root, train=False, transform=uniform, download=True)
    test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=False, num_workers=1, pin_memory=True)

    s_test_losses = []
    s_test_accs = []

    for ii in range(repeats):
        model = SiCNN_3(f_in, size, ratio, nratio, srange, padding)
        model.to(device)

        for epoch in range(1, nb_epochs + 1):
            train_l, train_a = train(model, train_loader, learning_rate, criterion, epoch, batch_log, device)
            train_l, train_a = test(model, train_loader, criterion, epoch, batch_log, device)

        test_l, test_a = test(model, test_loader, criterion, epoch, batch_log, device)

        s_test_losses.append(test_l)
        s_test_accs.append(test_a)

        pickle.dump(model, log)

    dynamics = {
        "scale": scale,
        "test_losses": s_test_losses,
        "test_accs": s_test_accs
    }
    pickle.dump(dynamics, log)

    mean_l = np.mean(np.array(s_test_losses))
    std_l = np.std(np.array(s_test_losses))
    mean_a = np.mean(np.array(s_test_accs))
    std_a = np.std(np.array(s_test_accs))

    avg_test_losses.append(mean_l)
    avg_test_accs.append(mean_a)
    std_test_losses.append(std_l)
    std_test_accs.append(std_a)

pickle.dump(avg_test_losses, log)
pickle.dump(std_test_losses, log)
pickle.dump(avg_test_accs, log)
pickle.dump(std_test_accs, log)

log.close()

plt.figure()
plt.errorbar([str(s) for s in scales], avg_test_losses, yerr=std_test_losses)
plt.title("Average loss vs Scale factor")
plt.xlabel("Scale range")
plt.ylabel("Categorical cross entropy")
plt.legend()
plt.savefig("test_loss_range_mean_sr0.pdf")

plt.figure()
plt.errorbar([str(s) for s in scales], avg_test_accs, yerr=std_test_accs)
plt.title("Average accuracy vs Scale factor")
plt.xlabel("Scale range")
plt.ylabel("Accuracy %")
plt.legend()
plt.savefig("test_acc_range_mean_sr0.pdf")

plt.figure()
plt.errorbar([str(s) for s in scales], [100-x for x in avg_test_accs], yerr=std_test_accs)
plt.title("Average error vs Test scale")
plt.xlabel("Test scale")
plt.ylabel("Error %")
plt.legend()
plt.savefig("test_err_range_mean_sr0.pdf")


import sys
from client_simulation import *
from message import *
import enum
import sip
from PyQt5.QtWidgets import QApplication, QWidget
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import Qt

# Use the more modern PyQt API (not enabled by default in Python 2.x);
# must precede importing any module that provides the API specified
sip.setapi('QDate', 2)
sip.setapi('QDateTime', 2)
sip.setapi('QString', 2)
sip.setapi('QTextStream', 2)
sip.setapi('QTime', 2)
sip.setapi('QUrl', 2)
sip.setapi('QVariant', 2)

# Default motor speeds {DRIVE_MOTORS, ACTUATOR, BUCKET}
MOTOR_SPEEDS = {0: 100, 1: 100, 2: 100}

MAX_MOTOR_SPEED = 100
BUFFER = []

class Motor(enum.Enum):
    DRIVE_MOTORS = 0
    ACTUATOR = 1
    BUCKET = 2

class Window(QWidget):
    def __init__(self, client):
        super(Window, self).__init__()

        self.client = client

        self.drive_keys_pressed = []

        self.actuator_keys_pressed = []

        self.bucket_keys_pressed = []

        self.motor_speed_to_adjust = Motor.DRIVE_MOTORS.value

        self.init_ui()

    def init_ui(self):
        # Button to open the connection
        self.open_connection_button = QtWidgets.QPushButton('Open Connection', self)
        self.open_connection_button.clicked.connect(self.open_connection_event)
        self.open_connection_button.setFocusPolicy(QtCore.Qt.NoFocus)

        # Button to close the connection
        self.close_connection_button = QtWidgets.QPushButton('Close Connection', self)
        self.close_connection_button.clicked.connect(self.close_connection_event)
        self.close_connection_button.setFocusPolicy(QtCore.Qt.NoFocus)
        self.close_connection_button.setEnabled(False)

        # Button to activate autonomy
        self.activate_autonomy_button = QtWidgets.QPushButton('Activate Autonomy', self)
        self.activate_autonomy_button.clicked.connect(self.activate_autonomy_event)
        self.activate_autonomy_button.setFocusPolicy(QtCore.Qt.NoFocus)
        self.activate_autonomy_button.setEnabled(False)

        # Button to deactivate autonomy
        self.deactivate_autonomy_button = QtWidgets.QPushButton('Deactivate Autonomy', self)
        self.deactivate_autonomy_button.clicked.connect(self.deactivate_autonomy_event)
        self.deactivate_autonomy_button.setFocusPolicy(QtCore.Qt.NoFocus)
        self.deactivate_autonomy_button.setEnabled(False)

        # Button to run simulation
        self.simulation_button = QtWidgets.QPushButton('Simulation', self)
        self.simulation_button.move(0, 30)
        self.simulation_button.clicked.connect(self.simulation_event)
        self.simulation_button.setFocusPolicy(QtCore.Qt.NoFocus)
        self.simulation_button.setEnabled(True)

        # Button to run real time tracking
        self.real_time_tracking_button = QtWidgets.QPushButton('Real Time Tracking', self)
        self.real_time_tracking_button.clicked.connect(self.real_time_tracking_event)
        self.real_time_tracking_button.setFocusPolicy(QtCore.Qt.NoFocus)
        self.real_time_tracking_button.setEnabled(False)

        grid =  QtWidgets.QGridLayout()
        grid.setSpacing(10)

        h_box_0 = QtWidgets.QHBoxLayout()
        h_box_0.addStretch(1)
        h_box_0.addWidget(self.open_connection_button)
        h_box_0.addWidget(self.close_connection_button)

        h_box_1 = QtWidgets.QHBoxLayout()
        h_box_1.addStretch(1)
        h_box_1.addWidget(self.activate_autonomy_button)
        h_box_1.addWidget(self.deactivate_autonomy_button)

        v_box = QtWidgets.QVBoxLayout()
        v_box.addStretch(1)
        v_box.addLayout(grid)
        v_box.addLayout(h_box_0)
        v_box.addLayout(h_box_1)

        self.setLayout(v_box)

        self.resize(500, 400)
        self.center()

        self.setWindowTitle('Client')
        self.setWindowIcon(QtGui.QIcon('mars.png'))

        self.show()

    def center(self):
        qr = self.frameGeometry()
        cp = QtWidgets.QDesktopWidget().availableGeometry().center()
        qr.moveCenter(cp)
        self.move(qr.topLeft())

    def keyPressEvent(self, event):
        if self.open_connection_button.isEnabled() \
                or self.deactivate_autonomy_button.isEnabled():
            return

        key = event.key()

        if not event.isAutoRepeat():
            # Driving logic
            if key == QtCore.Qt.Key_W:
                #BUFFER.append('W\n')
                if key not in self.drive_keys_pressed:
                    self.drive_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value],
                                'r': MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_S:
                #BUFFER.append('S\n')
                if key not in self.drive_keys_pressed:
                    self.drive_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': -1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value],
                                'r': -1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_A:
                #BUFFER.append('A\n')
                if key not in self.drive_keys_pressed:
                    self.drive_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': -1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value],
                                'r': 1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_D:
                #BUFFER.append('D\n')
                if key not in self.drive_keys_pressed:
                    self.drive_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': 1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value],
                                'r': -1 * MOTOR_SPEEDS[Motor.DRIVE_MOTORS.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

            # Motor speed adjustment mode logic
            elif key == QtCore.Qt.Key_1:
                #BUFFER.append('1\n')
                self.motor_speed_to_adjust = Motor.DRIVE_MOTORS.value
                print ('Motor speed adjustment mode:'), str(self.motor_speed_to_adjust)
            elif key == QtCore.Qt.Key_2:
                #BUFFER.append('2\n')
                self.motor_speed_to_adjust = Motor.ACTUATOR.value
                print ('Motor speed adjustment mode:'), str(self.motor_speed_to_adjust)
            elif key == QtCore.Qt.Key_3:
                #BUFFER.append('3\n')
                self.motor_speed_to_adjust = Motor.BUCKET.value
                print ('Motor speed adjustment mode:'), str(self.motor_speed_to_adjust)

            # Actuator logic
            elif key == QtCore.Qt.Key_U:
                #BUFFER.append('U\n')
                if key not in self.actuator_keys_pressed:
                    self.actuator_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'a': MOTOR_SPEEDS[Motor.ACTUATOR.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_J:
                #BUFFER.append('J\n')
                if key not in self.actuator_keys_pressed:
                    self.actuator_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'a': -1 * MOTOR_SPEEDS[Motor.ACTUATOR.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

            # Bucket logic
            elif key == QtCore.Qt.Key_I:
                #BUFFER.append('I\n')
                if key not in self.bucket_keys_pressed:
                    self.bucket_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'b': MOTOR_SPEEDS[Motor.BUCKET.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_K:
                #BUFFER.append('K\n')
                if key not in self.bucket_keys_pressed:
                    self.bucket_keys_pressed.append(key)

                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'b': -1 * MOTOR_SPEEDS[Motor.BUCKET.value]}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

        # Motor speed adjustment logic
        if key == QtCore.Qt.Key_Up:
            #BUFFER.append('Up\n')
            if MOTOR_SPEEDS[self.motor_speed_to_adjust] < MAX_MOTOR_SPEED:
                MOTOR_SPEEDS[self.motor_speed_to_adjust] += 1
                print (MOTOR_SPEEDS[self.motor_speed_to_adjust])
                if len(self.drive_keys_pressed):
                    self.keyPressEvent(Qt.QKeyEvent(Qt.QEvent.KeyPress, self.drive_keys_pressed[-1],
                                                    QtCore.Qt.NoModifier))
        elif key == QtCore.Qt.Key_Down:
            #BUFFER.append('Down\n')
            if MOTOR_SPEEDS[self.motor_speed_to_adjust] > 0:
                MOTOR_SPEEDS[self.motor_speed_to_adjust] -= 1
                print (MOTOR_SPEEDS[self.motor_speed_to_adjust])
                if len(self.drive_keys_pressed):
                    self.keyPressEvent(Qt.QKeyEvent(Qt.QEvent.KeyPress, self.drive_keys_pressed[-1],
                                                    QtCore.Qt.NoModifier))
    def keyReleaseEvent(self, event):
        if event.isAutoRepeat() or self.open_connection_button.isEnabled() \
                or self.deactivate_autonomy_button.isEnabled():
            return

        key = event.key()

        if key in self.drive_keys_pressed:
            self.drive_keys_pressed.remove(key)

        if key in self.actuator_keys_pressed:
            self.actuator_keys_pressed.remove(key)

        if key in self.bucket_keys_pressed:
            self.bucket_keys_pressed.remove(key)

        # Driving logic
        if not len(self.drive_keys_pressed):
            if key == QtCore.Qt.Key_W:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': 0, 'r': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_S:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': 0, 'r': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_A:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': 0, 'r': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_D:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'l': 0, 'r': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

        # Actuator logic
        if not len(self.actuator_keys_pressed):
            if key == QtCore.Qt.Key_U:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'a': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_J:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'a': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

        # Bucket logic
        if not len(self.bucket_keys_pressed):
            if key == QtCore.Qt.Key_I:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'b': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)
            elif key == QtCore.Qt.Key_K:
                forwarding_prefix = ForwardingPrefix.MOTOR.value
                sub_messages = {'b': 0}
                message = Message(forwarding_prefix, sub_messages).message
                self.client.send_message(message)

    def open_connection_event(self):
        self.client.open_connection()
        self.open_connection_button.setEnabled(False)
        self.close_connection_button.setEnabled(True)
        self.real_time_tracking_button.setEnabled(True)
        self.simulation_button.setEnabled(False)
        self.activate_autonomy_button.setEnabled(True)

    def close_connection_event(self):
        if self.deactivate_autonomy_button.isEnabled():
            autonomy_message = 'Please deactivate autonomy first.'
            QtWidgets.QMessageBox.question(self, 'Message',
                                           autonomy_message, QtWidgets.QMessageBox.Ok)
            return

        quit_message = 'Are you sure you want to close the connection?'
        reply = QtWidgets.QMessageBox.question(self, 'Message',
                                               quit_message,
                                               QtWidgets.QMessageBox.No,
                                               QtWidgets.QMessageBox.Yes)

        if reply == QtWidgets.QMessageBox.Yes:
            self.client.close_connection()

            self.open_connection_button.setEnabled(True)
            self.close_connection_button.setEnabled(False)

            self.activate_autonomy_button.setEnabled(False)
            self.deactivate_autonomy_button.setEnabled(False)

    def activate_autonomy_event(self):
        quit_message = 'Are you sure you want to activate autonomy?'
        reply = QtWidgets.QMessageBox.question(self, 'Message',
                                               quit_message,
                                               QtWidgets.QMessageBox.No,
                                               QtWidgets.QMessageBox.Yes)

        if reply == QtWidgets.QMessageBox.Yes:
            self.activate_autonomy()

    def activate_autonomy(self):
        self.client.send_message(ForwardingPrefix.CONTROLLER.value + AUTONOMY_ACTIVATION_MESSAGE)

        self.activate_autonomy_button.setEnabled(False)
        self.deactivate_autonomy_button.setEnabled(True)

    def deactivate_autonomy_event(self):
        deactivateMessage = 'Are you sure you want to deactivate autonomy?'
        reply = QtWidgets.QMessageBox.question(self, 'Message',
                                               deactivateMessage,
                                               QtWidgets.QMessageBox.No,
                                               QtWidgets.QMessageBox.Yes)

        if reply == QtWidgets.QMessageBox.Yes:
            self.deactivate_autonomy()

    def deactivate_autonomy(self):
        self.client.send_message(ForwardingPrefix.CONTROLLER.value + AUTONOMY_DEACTIVATION_MESSAGE)

        self.activate_autonomy_button.setEnabled(True)
        self.deactivate_autonomy_button.setEnabled(False)

    def simulation_event(self):
        sim = Setup()
        sim.start()

    def real_time_tracking_event(self):
        print()

    def closeEvent(self, QCloseEvent):
        quit_message = 'Are you sure you want to exit the client?'
        reply = QtWidgets.QMessageBox.question(self, 'Message',
                                               quit_message,
                                               QtWidgets.QMessageBox.No,
                                               QtWidgets.QMessageBox.Yes)

        if reply == QtWidgets.QMessageBox.Yes:
            if self.deactivate_autonomy_button.isEnabled():
                self.deactivate_autonomy()
            self.client.shutdown()
            QCloseEvent.accept()
        else:
            QCloseEvent.ignore()

def main(client):
    app = QApplication(sys.argv)
    window = Window(client)
    # log = open('key_press_log.txt', 'w')
    # log.write('Key Press Log: \n')
    # for i in BUFFER:
    #     log.write(i)
    # log.close()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


import os
import pickle
from itertools import islice

import torch
import torch.distributed as dist

from onnxruntime import set_seed
from onnxruntime.training import amp, checkpoint, optim, orttrainer
from orttraining_test_orttrainer_frontend import _load_pytorch_transformer_model
from onnxruntime.capi._pybind_state import set_cuda_device_id, get_mpi_context_world_rank, get_mpi_context_world_size

from _test_commons import generate_dummy_optim_state

from numpy.testing import assert_allclose, assert_array_equal

global_fp16_fp32_atol = 1e-3

def _train(trainer, train_data, batcher_fn, total_batch_steps = 5, seed = 1):
    """Runs train_step total_batch_steps number of times on the given trainer"""
    for i in range(total_batch_steps):
        torch.manual_seed(seed)
        set_seed(seed)
        data, targets = batcher_fn(train_data, i*35)
        trainer.train_step(data, targets)

def makedir(checkpoint_dir):
    """Creates a directory if checkpoint_dir does not exist"""
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir, exist_ok = True)

def _save(trainer, checkpoint_dir, state_dict_key_name):
    """Saves the ORTTrainer checkpoint and the complete state dictionary to the given checkpoint_dir directory"""

    # save current model parameters as a checkpoint
    makedir(checkpoint_dir)
    checkpoint.experimental_save_checkpoint(trainer, checkpoint_dir)
    state_dict = checkpoint.experimental_state_dict(trainer)
    pickle.dump({state_dict_key_name : state_dict}, open(os.path.join(checkpoint_dir, state_dict_key_name+'.pkl'), "wb"))

def _chunkify(sequence, num_chunks):
    """Breaks down a given sequence into num_chunks chunks"""
    quo, rem = divmod(len(sequence), num_chunks)
    return (sequence[i * quo + min(i, rem):(i + 1) * quo + min(i + 1, rem)] for i in range(num_chunks))

def _setup_test_infra(world_rank, world_size):
    """distributed setup just for testing purposes"""
    os.environ['RANK'] = str(world_rank)
    os.environ['WORLD_SIZE'] = str(world_size)
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'

    set_cuda_device_id(world_rank)

    dist.init_process_group(backend='nccl', world_size=world_size, rank=world_rank)

def distributed_setup(func):
    """Decorator function for distributed tests.

    Sets up distributed environment by extracting the following variables from MPI context
    - world_rank
    - world_size
    - device

    Also sets up the infrastructure required for the distributed tests such as setting up the torch distributed initialization
    """
    def setup(checkpoint_dir):
        world_rank = get_mpi_context_world_rank()
        world_size = get_mpi_context_world_size()
        device = 'cuda:' + str(world_rank)

        _setup_test_infra(world_rank, world_size)

        func(world_rank, world_size, device, checkpoint_dir=checkpoint_dir)

    return setup

def create_orttrainer_and_load_checkpoint(device, trainer_opts, checkpoint_dir, use_lamb=True):
    """Instantiate and load checkpoint into trainer

    - Instantiates the ORTTrainer with given input trainer_opts configuration for a simple transformer model
    - Loads the checkpoint from directory checkpoint_dir into the trainer
    - Runs eval_step on the trainer so the trainer onnx graph is initialized
    - Returns the trainer state_dict and the pytorch model
    """
    seed = 1
    torch.manual_seed(seed)
    set_seed(seed)

    # PyTorch transformer model setup
    learning_rate = 0.1
    optim_config = optim.LambConfig(lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
    model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
    trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts))

    # load checkpoint into trainer
    checkpoint.experimental_load_checkpoint(trainer, checkpoint_dir)

    # run an eval step to innitialize the graph
    torch.manual_seed(seed)
    set_seed(seed)
    data, targets = batcher_fn(train_data, 0)
    trainer.eval_step(data, targets)

    return checkpoint.experimental_state_dict(trainer), model

def create_initialized_orttrainer(device, trainer_opts, use_lamb=True):
    seed = 1
    torch.manual_seed(seed)
    set_seed(seed)

    learning_rate = 1e-10
    optim_config = optim.LambConfig(lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
    model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
    trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts))

    _train(trainer, train_data, batcher_fn)

    return trainer

def verify_model_state(trainer, expected_state_dict, is_mixedprecision):
    actual_model_state = trainer._training_session.get_model_state(include_mixed_precision_weights=True)
    for fp_or_mp, value in actual_model_state.items():
        for weight_name in value:
            assert weight_name.find('_view_') == -1
    assert len(expected_state_dict['fp32_param']) == len(actual_model_state['full_precision']), \
        "expected and actual should have same number of tensors"
    for weight_name, tensor in expected_state_dict['fp32_param'].items():
        if not weight_name in actual_model_state['full_precision']:
            assert '_view_' in weight_name, \
                "only zero shared weight may not match name"
            weight_name = weight_name.split('_view_')[0]
        assert_allclose(tensor, actual_model_state['full_precision'][weight_name])

    if is_mixedprecision:
        assert 'mixed_precision' in actual_model_state.keys(), "missing 'mixed_precision' key in mixed precision run"
        assert len(expected_state_dict['fp16_param']) == len(actual_model_state['mixed_precision']), \
            "expected and actual should have same number of tensors"
        for weight_name, tensor in expected_state_dict['fp16_param'].items():
            weight_name = weight_name.split('_fp16')[0]
            assert_allclose(tensor, actual_model_state['mixed_precision'][weight_name])

def verify_opt_state(trainer, expected_state_dict):
    actual_opt_state = trainer._training_session.get_optimizer_state()
    actual_opt_count = sum(len(v) for v in actual_opt_state.values())
    assert actual_opt_count == len(expected_state_dict['optimizer'])
    for weight_name in actual_opt_state:
        assert weight_name.find('_view_') == -1
    for opt_name, expected_tensor in expected_state_dict['optimizer'].items():
        if opt_name == "Step":
            actual_tensor = actual_opt_state['shared_optimizer_state']['Step']
        else:
            if opt_name.startswith('Moment_'):
                prefix = opt_name[:len("Moment_0")]
                weight_name = opt_name[len("Moment_0_"):]
                if not weight_name in actual_opt_state:
                    assert '_view_' in weight_name, \
                        "only zero shared weight may not match name"
                    weight_name = weight_name.split('_view_')[0]
            elif opt_name.startswith('Update_Count_'):
                prefix = "Update_Count"
                weight_name = opt_name[len(prefix + 1):]
            actual_tensor = actual_opt_state[weight_name][prefix]
        assert_allclose(actual_tensor, expected_tensor, atol=global_fp16_fp32_atol)

def verify_part_info(trainer, expected_state_dict, is_mixedprecision, is_zero_run):
    part_info = trainer._training_session.get_partition_info_map()
    for weight_name, weight_info in part_info.items():
        for info, value in weight_info.items():
            assert isinstance(value, list), "get_partition_info_map should return list"
            assert isinstance(value[0], int), "get_partition_info_map should return list of int"
            if info == "megatron_row_partition":
                assert len(value) == 1, "megatron_row_partition should only have 1 element"
                if is_zero_run:
                    assert value[0] == -1, "megatron_row_partition is -1 if megatron optimization is not on"
            if info == "original_dim":
                if is_zero_run:
                    assert len(value) > 0, "original_dim should not be empty in zero run"
                    if is_mixedprecision:
                        assert_array_equal(part_info[weight_name]['original_dim'], expected_state_dict['fp16_param'][weight_name + '_fp16'].shape)
                    else:
                        assert_array_equal(part_info[weight_name]['original_dim'], expected_state_dict['fp32_param'][weight_name].shape)

def split_state_dict(state_dict):
    """Given a flat state dictionary, split it into optimizer, fp32_param, fp16_param hierarchical dictionary and return"""

    optimizer_keys = ['Moment_1_', 'Moment_2_', 'Update_Count_', 'Step']
    split_sd = {'optimizer': {}, 'fp32_param': {}, 'fp16_param': {}}
    for k, v in state_dict.items():
        mode = 'fp32_param'
        for optim_key in optimizer_keys:
            if k.startswith(optim_key):
                mode = 'optimizer'
                break
        if k.endswith('_fp16'):
            mode = 'fp16_param'
        split_sd[mode][k] = v
    return split_sd

def _split_name(name):
    """Splits given state name (model or optimizer state name) into the param_name, optimizer_key, view_num and the fp16_key"""
    name_split = name.split('_view_')
    view_num = None
    if(len(name_split) > 1):
        view_num = int(name_split[1])
    optimizer_key = ''
    fp16_key = ''
    if name_split[0].startswith('Moment_1'):
        optimizer_key = 'Moment_1_'
    elif name_split[0].startswith('Moment_2'):
        optimizer_key = 'Moment_2_'
    elif name_split[0].startswith('Update_Count'):
        optimizer_key = 'Update_Count_'
    elif name_split[0].endswith('_fp16'):
        fp16_key = '_fp16'
    param_name = name_split[0]
    if optimizer_key != '':
        param_name = param_name.split(optimizer_key)[1]
    param_name = param_name.split('_fp16')[0]
    return param_name, optimizer_key, view_num, fp16_key

def aggregate_states(aggregated_states, state_dict):
    """Concatenate existing aggregated state dict values with given state_dict values"""

    for key, value in state_dict.items():
        weight_name, optimizer_key, view_num, fp16_key = _split_name(key)
        if view_num is not None:
            # parameter is sharded
            param_name = optimizer_key + weight_name + fp16_key

            if param_name in aggregated_states and optimizer_key not in ['Update_Count_']:
                # found a previous shard of the param, concatenate shards ordered by ranks
                aggregated_states[param_name] = torch.cat((aggregated_states[param_name], value))
            else:
                aggregated_states[param_name] = value
        else:
            aggregated_states[key] = value

def create_orttrainer_and_save_checkpoint(device, trainer_opts, checkpoint_dir, state_dict_key_name='state_dict', use_lamb=True):
    learning_rate = 0.1
    seed = 1

    torch.manual_seed(seed)
    set_seed(seed)

    optim_config = optim.LambConfig(lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
    model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
    trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts))

    if 'distributed' in trainer_opts:
        train_data = next(islice(_chunkify(train_data, trainer_opts['distributed']['world_size']), trainer_opts['distributed']['world_rank'], None))

    # run train steps
    _train(trainer, train_data, batcher_fn)

    # save current model parameters as a checkpoint
    if checkpoint_dir:
        _save(trainer, checkpoint_dir, state_dict_key_name)

def load_model_optim_state_and_eval(device, trainer_opts, use_lamb=True):
    learning_rate = 0.1
    seed = 1

    torch.manual_seed(seed)
    set_seed(seed)

    optim_config = optim.LambConfig(lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
    model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
    trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts))

    # load dummy state
    dummy_init_state = generate_dummy_optim_state(model, optim_config)
    checkpoint._experimental_load_optimizer_state(trainer, dummy_init_state)

    # run an eval step to innitialize the graph
    data, targets = batcher_fn(train_data, 0)
    trainer.eval_step(data, targets)

    return dummy_init_state, checkpoint.experimental_state_dict(trainer)


# -*- coding: utf-8 -*-
import sqlite3
import base64
import requests
import json
import hashlib
import logging
from multiprocessing.util import register_after_fork
import os
import base64
import hashlib
import shutil
import transaction
import tempfile
from pydub import AudioSegment
from collections import defaultdict
from pathvalidate import sanitize_filename
from urllib import request

from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy import create_engine
from sqlalchemy import and_
from lingvodoc.models import (
    Client,
    DBSession as SyncDBSession,
    UserBlobs,
    TranslationAtom,
    TranslationGist,
    Field,
    Entity,
    LexicalEntry,
    Dictionary,
    Language,
    User,
    DictionaryPerspectiveToField,
    DictionaryPerspective,
    BaseGroup,
    Group,
    PublishingEntity

)
from lingvodoc.cache.caching import TaskStatus, initialize_cache

from sqlalchemy.orm import (
    sessionmaker,
)

from sqlalchemy import tuple_, and_, or_
from lingvodoc.scripts import elan_parser
from pdb import set_trace
import io
import xlsxwriter
# from time import time
import datetime
import time
import traceback
from os import path, makedirs
from errno import EEXIST
from shutil import copyfileobj

EAF_TIERS = {
    "literary translation": "Translation of Paradigmatic forms",
    "text": "Transcription of Paradigmatic forms",
    "word": "Word",
    "transcription": "Transcription",
    "translation": "Translation"
}
log = logging.getLogger(__name__)
log.setLevel(logging.DEBUG)
field_ids_to_str = lambda x: str(x.field_client_id) + '_' + str(x.field_object_id)

def is_empty(row):
    for col in row:
        if col:
            return False
    return True

def find_lexical_entries_by_tags(tags, session, published):
    result = session.query(LexicalEntry) \
        .join(Entity) \
        .join(PublishingEntity) \
        .filter(Entity.content.in_(tags),
                PublishingEntity.accepted == True,
                Entity.field_client_id == 66,
                Entity.field_object_id == 25)
    if published is not None:
        result = result.filter(PublishingEntity.published == published)
    return result.all()

def find_all_tags(tag, session, published):
    tags = [tag]
    new_tags = [tag]
    while new_tags:
        lexical_entries = find_lexical_entries_by_tags(new_tags, session, published)
        new_tags = list()
        for lex in lexical_entries:
            entities = session.query(Entity)  \
                .join(PublishingEntity) \
                .filter(Entity.parent == lex,
                        PublishingEntity.accepted == True,
                        Entity.field_client_id == 66,
                        Entity.field_object_id == 25)
            if published is not None:
                entities = entities.filter(PublishingEntity.published == published)
            for entity in entities:
                if entity.content not in tags:
                    tags.append(entity.content)
                    new_tags.append(entity.content)
    return tags

def generate_massive_cell(tag, session, text_fields, published):
    result = list()
    tags = find_all_tags(tag, session, published)
    lexical_entries = find_lexical_entries_by_tags(tags, session, published)
    for lex in lexical_entries:
        entities = session.query(Entity).join(PublishingEntity) \
            .filter(Entity.parent_client_id == lex.client_id,
                    Entity.parent_object_id == lex.object_id,
                    PublishingEntity.accepted == True,
                    Entity.marked_for_deletion == False)
        if published is not None:
            entities = entities.filter(PublishingEntity.published == published)
        subres = []
        for entity in entities:
            if (entity.field_client_id, entity.field_object_id) in text_fields and entity.content is not None:
                subres.append(entity.content)
        if len(subres) > 0:
            result.append("; ".join(subres))
    if len(result) > 0:
        return "\n".join(result)
    else:
        return ""

def compile_workbook(client_id, object_id, workbook_stream, session, locale_id, published):
    """
    Compiles analysis results into an Excel workbook.
    """

    workbook = xlsxwriter.Workbook(workbook_stream, {'in_memory': True})

    dictionary = session.query(Dictionary).filter_by(client_id=client_id, object_id=object_id,
                                                     marked_for_deletion=False).one()

    text_fields = session.query(Field.client_id, Field.object_id) \
        .filter_by(data_type_translation_gist_client_id=1,
                   data_type_translation_gist_object_id=47).all()
    perspectives = session.query(DictionaryPerspective).filter_by(parent_client_id=client_id,
                                                                  parent_object_id=object_id,
                                                                  marked_for_deletion=False).all()
    for perspective in perspectives:
        perspective_name = perspective.get_translation(client_id, session) + "_" + str(
            perspective.client_id) + "_" + str(perspective.object_id)
        for c in u"[]:?/*\x00":
            perspective_name = perspective_name.replace(c, "")
        if len(perspective_name) >= 31:
            perspective_name = perspective_name[:20] + "_" + str(perspective.client_id) + "_" + str(
                perspective.object_id)
        if len(perspective_name) >= 31:
            perspective_name = perspective_name[:10] + "_" + str(perspective.client_id) + "_" + str(
                perspective.object_id)
        worksheet = workbook.add_worksheet(name=perspective_name)
        fields = session.query(DictionaryPerspectiveToField).filter_by(parent_client_id=perspective.client_id,
                                                                       parent_object_id=perspective.object_id,
                                                                       marked_for_deletion=False,
                                                                       ).filter(
            tuple_(DictionaryPerspectiveToField.field_client_id, DictionaryPerspectiveToField.field_object_id).in_(
                text_fields)
        ).order_by(DictionaryPerspectiveToField.position).all()
        etymology_field = session.query(DictionaryPerspectiveToField).filter_by(parent_client_id=perspective.client_id,
                                                                                parent_object_id=perspective.object_id,
                                                                                marked_for_deletion=False,
                                                                                field_client_id=66,
                                                                                field_object_id=25
                                                                                ).first()

        field_to_column = {field_ids_to_str(field): counter for counter, field in
                           enumerate(fields)}
        row = 1
        column = 0
        for field in fields:
            worksheet.write(0, column, field.field.get_translation(locale_id, session))
            column += 1

        if etymology_field:
            worksheet.write(0, column, etymology_field.field.get_translation(locale_id, session))

        lexical_entries = session.query(LexicalEntry).join(Entity).join(PublishingEntity) \
            .filter(LexicalEntry.parent_client_id == perspective.client_id,
                    LexicalEntry.parent_object_id == perspective.object_id,
                    LexicalEntry.marked_for_deletion == False,
                    Entity.marked_for_deletion == False,
                    PublishingEntity.accepted == True)
        if published is not None:
            lexical_entries = lexical_entries.filter(PublishingEntity.published == published)
        for lex in lexical_entries:
            row_to_write = ["" for field in fields]
            entities = session.query(Entity).join(PublishingEntity) \
                .filter(Entity.parent_client_id == lex.client_id,
                        Entity.parent_object_id == lex.object_id,
                        Entity.marked_for_deletion == False,
                        PublishingEntity.accepted == True)
            if published is not None:
                entities = entities.filter(PublishingEntity.published == published)
            for entity in entities:
                ent_field_ids = field_ids_to_str(entity)
                if ent_field_ids in field_to_column:
                    if row_to_write[field_to_column[ent_field_ids]] == "":
                        row_to_write[field_to_column[ent_field_ids]] = entity.content
                    else:
                        row_to_write[field_to_column[ent_field_ids]] += "\n" + entity.content
                if etymology_field and len(row_to_write) == len(fields) and ent_field_ids == "66_25":
                    row_to_write.append(generate_massive_cell(entity.content, session, text_fields, published))
            if not is_empty(row_to_write):
                worksheet.write_row(row, 0, row_to_write)
                row += 1
    return

# @profile()
def save(
        client_id,
        object_id,
        storage,
        sqlalchemy_url,
        task_key,
        cache_kwargs,
        dict_name,
        locale_id,
        published
):  # :(

    initialize_cache(cache_kwargs)
    task_status = TaskStatus.get_from_cache(task_key)

    engine = create_engine(sqlalchemy_url)
    register_after_fork(engine, engine.dispose)
    log = logging.getLogger(__name__)
    Session = sessionmaker(bind=engine)
    session = Session()
    task_status.set(3, 20, 'Running async process')

    workbook_stream = io.BytesIO()

    try:
        compile_workbook(client_id, object_id, workbook_stream, session, locale_id, published)

        workbook_stream.seek(0)

    except Exception as exception:

        traceback_string = ''.join(traceback.format_exception(
            exception, exception, exception.__traceback__))[:-1]

        log.debug('compile_workbook: exception')
        log.debug(traceback_string)

        # If we failed to create an Excel file, we terminate with error.

        task_status.set(4, 100,
                        'Finished (ERROR), result compilation error')

        return {'error': 'result compilation error'}

    # Name(s) of the resulting file(s) includes dictionary name, perspective name and current date.

    current_datetime = datetime.datetime.now(datetime.timezone.utc)

    result_filename = '{0} - {1:04d}.{2:02d}.{3:02d}'.format(
        dict_name[:64],
        current_datetime.year,
        current_datetime.month,
        current_datetime.day)

    table_filename = sanitize_filename(result_filename + '.xlsx')

    # cur_time = time.time()
    dictionary = session.query(Dictionary).filter_by(client_id=client_id, object_id=object_id).one()
    dict_status_atom = session.query(TranslationAtom).filter_by(
        parent_client_id=dictionary.state_translation_gist_client_id,
        parent_object_id=dictionary.state_translation_gist_object_id,
        locale_id=2).first()
    if not dict_status_atom:
        dict_status = 'translation_failure'
    else:
        dict_status = dict_status_atom.content

    if published is None:
        cur_folder = 'edit'
    elif published is True:
        cur_folder = 'view'
    else:
        cur_folder = 'should_be_impossible'

    storage_dir = path.join(storage['path'], 'save_dictionary', dict_status, cur_folder)
    makedirs(storage_dir, exist_ok=True)

    # Storing file with the results.

    storage_path = path.join(storage_dir, table_filename)
    directory = path.dirname(storage_path)

    try:
        makedirs(directory)

    except OSError as exception:
        if exception.errno != EEXIST:
            raise

    # If the name of the result file is too long, we try again with a shorter name.

    try:
        with open(storage_path, 'wb+') as workbook_file:
            copyfileobj(workbook_stream, workbook_file)

    except OSError as os_error:

        if os_error.errno != 36:
            raise

        result_filename = '{0} - {1:04d}.{2:02d}.{3:02d}'.format(
            dict_name[:32],
            current_datetime.year,
            current_datetime.month,
            current_datetime.day)

        table_filename = sanitize_filename(result_filename + '.xlsx')
        storage_path = path.join(storage_dir, table_filename)

        with open(storage_path, 'wb+') as workbook_file:
            copyfileobj(workbook_stream, workbook_file)

    # Successfully compiled phonology, finishing and returning links to files with results.

    url_list = [

        ''.join([
            storage['prefix'],
            storage['static_route'],
            'save_dictionary', '/',
            dict_status, '/',
            cur_folder, '/',
            filename])

        for filename in [table_filename]]

    task_status.set(4, 100, 'Finished', result_link_list=url_list)

    session.commit()
    engine.dispose()
    return

def save_dictionary(
        client_id,
        object_id,
        storage,
        sqlalchemy_url,
        task_key,
        cache_kwargs,
        dict_name,
        locale_id,
        published
):
    save(
        client_id,
        object_id,
        storage,
        sqlalchemy_url,
        task_key,
        cache_kwargs,
        dict_name,
        locale_id,
        published
    )


import os
import torch
import numpy as np
from PIL import Image
import math

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

createImg_root = '/mnt/valid_rotate180/images'
createLab_root = '/mnt/valid_rotate180/labels'

def get_rotate_mat(theta):
    return np.array([[math.cos(theta), -math.sin(theta)], [math.sin(theta), math.cos(theta)]])

def rotate_vertices(vertices, theta, anchor=None):
    v = vertices.reshape((4,2)).T
    if anchor is None:
        anchor = v[:, :1]
    rotate_mat = get_rotate_mat(theta)
    res = np.dot(rotate_mat, v - anchor)
    return (res + anchor).T.reshape(-1)

def rotate_img(img_dir, img_savepath, file_path, lab_savepath, flag):
    image = Image.open(img_dir)
    if flag:
        image = image.rotate(180, Image.BILINEAR)
    image.save(img_savepath)

    with open(lab_savepath, 'w') as writer:
        with open(file_path, 'r') as lines:
            lines = lines.readlines()
            for l, line in enumerate(lines):
                line = line.split(';')
                vertice = [int(vt) for vt in line[1:-1]]
                vertice = np.array(vertice)
                if flag:
                    center_x = (image.width - 1) / 2
                    center_y = (image.height - 1) / 2
                    new_vertice = np.zeros(vertice.shape)
                    new_vertice[:] = rotate_vertices(vertice, - math.pi, np.array([[center_x], [center_y]]))
                    vertice = new_vertice
                new_line = []
                new_line.append(line[0])
                for v in vertice:
                    new_line.append(str(int(v)))
                new_line.append(line[-1])
                new_line = ';'.join(new_line)
                writer.write(new_line)
        writer.close()

if __name__ == '__main__':

    img_label_dir = '/mnt/valid_rotate90/labels'
    for root, dirs, files in os.walk(img_label_dir):
        for file in sorted(files):
            file_path = os.path.join(root, file)
            image_name = file[0: -4] + '.jpg'
            image_dir = os.path.join('/mnt/valid_rotate90/images', image_name)
            with open(file_path, 'r') as lines:
                lines = lines.readlines()
                first_y = int(lines[0].split(';')[2])
                last_y = int(lines[0].split(';')[-2])
                flag = True
                if first_y < last_y:
                    flag = False
                img_savepath = os.path.join(createImg_root, image_name)
                lab_savepath = os.path.join(createLab_root, file)
                rotate_img(image_dir, img_savepath, file_path, lab_savepath, flag)


import os
import numpy as np
import tensorflow as tf
import cv2
import math
import time
import shutil
import cfgs
from tqdm import tqdm
from CenterNet import CenterNet
from utils.generator import get_data
from net.resnet import load_weights
import tensorflow.contrib.slim as slim

def train():
    # define dataset
    num_train_imgs = len(open(cfgs.TRAIN_DATA_FILE, 'r').readlines())
    num_train_batch = int(math.ceil(float(num_train_imgs) / cfgs.BATCH_SIZE))
    num_test_imgs = len(open(cfgs.TEST_DATA_FILE, 'r').readlines())
    num_test_batch = int(math.ceil(float(num_test_imgs) / 4))

    # train dataset
    train_dataset = tf.data.TextLineDataset(cfgs.TRAIN_DATA_FILE)
    train_dataset = train_dataset.shuffle(num_train_imgs)
    train_dataset = train_dataset.batch(cfgs.BATCH_SIZE)
    train_dataset = train_dataset.map(lambda x: tf.py_func(get_data, inp=[x, cfgs.USE_AUG],
                                                           Tout=[tf.float32, tf.float32, tf.float32, tf.float32,
                                                                 tf.float32, tf.float32, tf.float32]),
                                      num_parallel_calls=8)
    train_dataset = train_dataset.prefetch(8)

    # test dataset
    test_dataset = tf.data.TextLineDataset(cfgs.TEST_DATA_FILE)
    test_dataset = test_dataset.batch(4)
    test_dataset = test_dataset.map(lambda x: tf.py_func(get_data, inp=[x, False],
                                                         Tout=[tf.float32, tf.float32, tf.float32, tf.float32,
                                                               tf.float32, tf.float32, tf.float32]),
                                    num_parallel_calls=8)
    test_dataset = test_dataset.prefetch(8)

    iterator = tf.data.Iterator.from_structure(train_dataset.output_types, train_dataset.output_shapes)
    trainset_init_op = iterator.make_initializer(train_dataset)
    testset_init_op = iterator.make_initializer(test_dataset)

    input_data, batch_hm, batch_wh, batch_reg, batch_reg_mask, batch_ind, batch_cls = iterator.get_next()
    input_data.set_shape([cfgs.BATCH_SIZE, cfgs.INPUT_IMAGE_H, cfgs.INPUT_IMAGE_W, 3])
    batch_hm.set_shape(
        [cfgs.BATCH_SIZE, cfgs.INPUT_IMAGE_H // cfgs.DOWN_RATIO, cfgs.INPUT_IMAGE_W // cfgs.DOWN_RATIO, cfgs.NUM_CLASS])
    batch_wh.set_shape([cfgs.BATCH_SIZE, cfgs.MAX_OBJ, 2])
    batch_reg.set_shape([cfgs.BATCH_SIZE, cfgs.MAX_OBJ, 2])
    batch_reg_mask.set_shape([cfgs.BATCH_SIZE, cfgs.MAX_OBJ])
    batch_ind.set_shape([cfgs.BATCH_SIZE, cfgs.MAX_OBJ])
    batch_cls.set_shape([cfgs.BATCH_SIZE, cfgs.MAX_OBJ])

    # training flag
    is_training = tf.placeholder(dtype=tf.bool, name='is_training')

    # define model and loss
    model = CenterNet(input_data, is_training)
    with tf.variable_scope('loss'):
        hm_loss, wh_loss, reg_loss, cls_loss = model.compute_loss(batch_hm, batch_wh, batch_reg, batch_reg_mask,
                                                                  batch_ind, batch_cls)
        regular_loss = tf.add_n(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))
        total_loss = hm_loss + wh_loss + reg_loss + cls_loss + regular_loss

    # define train op
    if cfgs.LR_TPYE == "CosineAnnealing":
        global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
        warmup_steps = tf.constant(cfgs.WARM_UP_EPOCHS * num_train_batch, dtype=tf.float64, name='warmup_steps')
        train_steps = tf.constant(cfgs.EPOCHS * num_train_batch, dtype=tf.float64, name='train_steps')
        learning_rate = tf.cond(
            pred=global_step < warmup_steps,
            true_fn=lambda: global_step / warmup_steps * cfgs.INIT_LR,
            false_fn=lambda: cfgs.END_LR + 0.5 * (cfgs.INIT_LR - cfgs.END_LR) *
                             (1 + tf.cos(
                                 (global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
        )
        global_step_update = tf.assign_add(global_step, 1.0)

        optimizer = tf.train.AdamOptimizer(learning_rate).minimize(total_loss)
        with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
            with tf.control_dependencies([optimizer, global_step_update]):
                train_op = tf.no_op()

    else:
        global_step = tf.Variable(0, trainable=False)
        if cfgs.LR_TPYE == "exponential":
            learning_rate = tf.train.exponential_decay(cfgs.LR,
                                                       global_step,
                                                       cfgs.LR_DECAY_STEPS,
                                                       cfgs.LR_DECAY_RATE,
                                                       staircase=True)
        elif cfgs.LR_TPYE == "piecewise":
            learning_rate = tf.train.piecewise_constant(global_step, cfgs.LR_BOUNDARIES, cfgs.LR_PIECEWISE)
        optimizer = tf.train.AdamOptimizer(learning_rate)
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            train_op = optimizer.minimize(total_loss, global_step=global_step)

    saver = tf.train.Saver(tf.global_variables(), max_to_keep=cfgs.MAX_KEEP)
    pre_test_loss = float('inf')

    with tf.Session() as sess:
        with tf.name_scope('summary'):
            tf.summary.scalar("learning_rate", learning_rate)
            tf.summary.scalar("hm_loss", hm_loss)
            tf.summary.scalar("wh_loss", wh_loss)
            tf.summary.scalar("reg_loss", reg_loss)
            tf.summary.scalar('cls_loss', cls_loss)
            tf.summary.scalar('regular_loss', regular_loss)
            tf.summary.scalar("total_loss", total_loss)

            logdir = './log/' + cfgs.VERSION
            if os.path.exists(logdir): shutil.rmtree(logdir)
            os.mkdir(logdir)
            write_op = tf.summary.merge_all()
            summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph)

            ckptdir = './checkpoint/' + cfgs.VERSION
            if os.path.exists(ckptdir): shutil.rmtree(ckptdir)
            os.mkdir(ckptdir)

        # train
        sess.run(tf.global_variables_initializer())
        if cfgs.PRE_TRAIN:
            load_weights(sess, './pretrained_weights/mobilenet_v2.npy')
        for epoch in range(1, 1 + cfgs.EPOCHS):
            pbar = tqdm(range(num_train_batch))
            train_epoch_loss, test_epoch_loss = [], []
            train_hm_loss, train_wh_loss, train_reg_loss, train_cls_loss, train_regular_loss = [], [], [], [], []
            sess.run(trainset_init_op)
            for i in pbar:
                _, summary, train_step_loss, global_step_val, _hm_loss, _wh_loss, _reg_loss, _cls_loss, _regular_loss = sess.run(
                    [train_op, write_op, total_loss, global_step, hm_loss, wh_loss, reg_loss, cls_loss, regular_loss],
                    feed_dict={is_training: True})

                train_epoch_loss.append(train_step_loss)
                train_hm_loss.append(_hm_loss)
                train_wh_loss.append(_wh_loss)
                train_reg_loss.append(_reg_loss)
                train_cls_loss.append(_cls_loss)
                train_regular_loss.append(_regular_loss)
                summary_writer.add_summary(summary, global_step_val)
                pbar.set_description("train loss: %.2f" % train_step_loss)

            sess.run(testset_init_op)
            for j in range(num_test_batch):
                test_step_loss = sess.run(total_loss, feed_dict={is_training: False})
                test_epoch_loss.append(test_step_loss)

            train_epoch_loss, test_epoch_loss = np.mean(train_epoch_loss), np.mean(test_epoch_loss)
            train_hm_loss, train_wh_loss, train_reg_loss, train_cls_loss, train_regular_loss = np.mean(
                train_hm_loss), np.mean(train_wh_loss), np.mean(train_reg_loss), np.mean(train_cls_loss), np.mean(
                train_regular_loss)
            ckpt_file = ckptdir + "/centernet_test_loss=%.4f.ckpt" % test_epoch_loss
            log_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            print("=> Epoch: %2d Time: %s Train loss: %.2f Test loss: %.2f"
                  % (epoch, log_time, train_epoch_loss, test_epoch_loss))
            print('hm loss: %.4f  wh loss: %.4f  reg loss: %.4f  cls loss: %.4f  regular loss: %.4f' % (
                train_hm_loss, train_wh_loss, train_reg_loss, train_cls_loss, train_regular_loss))
            if cfgs.SAVE_MIN:
                if test_epoch_loss < pre_test_loss:
                    pre_test_loss = test_epoch_loss
                    print('Saving  %s' % (ckpt_file))
                    saver.save(sess, ckpt_file, global_step=epoch)
            else:
                saver.save(sess, ckpt_file, global_step=epoch)

if __name__ == '__main__': train()


# Copyright 2014, Sandia Corporation. Under the terms of Contract
# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains certain
# rights in this software.

from behave import *
import nose.tools

import glob
import logging
import os
import pkgutil
import shutil
import subprocess
import sys

import IPython
import IPython.core.interactiveshell
import nbformat

import toyplot

log = logging.getLogger(__name__)
log.name = "features.steps.documentation"

root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
docs_dir = os.path.join(root_dir, "docs")
package_dir = os.path.join(root_dir, "toyplot")

@given(u'all public Toyplot modules')
def step_impl(context):
    def walk_modules(package, path):
        modules = []
        modules.append(package)
        for loader, name, is_package in pkgutil.iter_modules([path]):
            modules += walk_modules(package + "." + name, os.path.join(path, name))
        return modules
    context.modules = sorted(walk_modules("toyplot", package_dir))

@given(u'the Toyplot reference documentation')
def step_impl(context):
    context.references = []
    for directory, subdirectories, filenames in os.walk(docs_dir):
        if os.path.basename(directory) in ["js"]:
            continue
        for filename in filenames:
            if os.path.splitext(filename)[1] not in [".rst"]:
                continue
            if not filename.startswith("toyplot."):
                continue

            context.references.append(os.path.join(directory, filename))
    context.references = sorted(context.references)

@then(u'every module must have a section in the reference documentation')
def step_impl(context):
    for module in context.modules:
        if os.path.join(docs_dir, module + ".rst") not in context.references:
            raise AssertionError("No matching documentation found for the %s module." % module)

@then(u'every section in the reference documentation must match a module')
def step_impl(context):
    modules = [os.path.join(docs_dir, module + ".rst") for module in context.modules]
    for reference in context.references:
        if reference not in modules:
            raise AssertionError("No matching module found for %s." % reference)

@given(u'the Toyplot documentation notebooks')
def step_impl(context):
    context.notebooks = sorted(glob.glob(os.path.join(docs_dir, "*.ipynb")))

@then(u'every notebook runs without error')
def step_impl(context):
    sys.path.append(docs_dir)
    cwd = os.getcwd()
    os.chdir(docs_dir)
    for notebook in context.notebooks:
        context.execute_steps(u"Then notebook %s runs without error" % notebook)
    os.chdir(cwd)
    sys.path.remove(docs_dir)

@then(u'notebook {notebook} runs without error')
def step_impl(context, notebook):
    log.info(notebook)

    with open(notebook) as stream:
        notebook = nbformat.read(stream, as_version=4)

    shell = IPython.core.interactiveshell.InteractiveShell.instance()
    nblocals = dict()

    for cell in notebook.cells:
        if cell.cell_type == "code":
            code = shell.input_transformer_manager.transform_cell(cell.source)
            exec(code, nblocals)
            toyplot.Canvas._ipython_post_execute()

#    # Run each notebook script, keeping track of failures.
#    failures = {}
#    for notebook, script in zip(context.notebooks, context.scripts):
#        try:
#            command = ["coverage", "run", "--source", "toyplot", "--append", script]
#            log.info(" ".join(command))
#            subprocess.check_call(command, cwd=docs_dir)
#            # Remove the script (we only do this if execution succeeded)
#            os.remove(script)
#        except Exception as e:
#            failures[notebook] = e
#
#    if failures:
#        message = ""
#        for notebook_path, exception in failures.items():
#            message += notebook_path + " exception:\n\n" + exception + "\n\n"
#        raise AssertionError(message)


import numpy as np
import cv2
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

video_capture = cv2.VideoCapture(0)

while True:
	ret, frame = video_capture.read()
	gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

	faces =face_cascade.detectMultiScale(gray, 1.3, 5)
	for (x,y,w,h) in faces:
		cv2.rectangle(frame, (x,y),(x+w,y+h),(0,255,0),2)
        # cv2.putText(frame,'PERSON',(50,50),cv2.FONT_HERSHEY_COMPLEX,1,(0,255,0),2)

	cv2.imshow('video',frame)
	if cv2.waitKey(1) & 0xFF == ord(' '):
		break
video_capture.release()
cv2.destroyAllWindows()


# coding:utf-8
from flask import Flask, jsonify, make_response, request, Response
from flask_cors import CORS
from yt_dlp import YoutubeDL

app = Flask(__name__)
CORS(app)

def downloadMedia(url, op):
	with YoutubeDL(op) as ydl:
		ydl.download([url])

def responseJSON(status, uID):
	if status == 200:
		status = "OK"
	else:
		status = "Error"
	response = {
			"data": {
				"status": status,
				"url": uID
				}
			}
	return response

# https://address/md/api/mp3?v=hoge
@app.route('/md/api/mp3', methods=['POST'])
def downloadMp3():
	options = {
			'format': 'bestaudio[ext=m4a]/bestaudio[ext=mp3]/bestaudio',
			'outtmpl': './Music/%(title)s.%(ext)s',
            'postprocessors':
            [{'key':'FFmpegExtractAudio','preferredcodec':'mp3','preferredquality':'192'}],
			}

	ID = request.args.get("v")
	if ID == None:
		return jsonify(responseJSON(0, "None"))
	else:
		URL = "https://www.youtube.com/watch?v="+str(ID)
		downloadMedia(URL, options)
		return jsonify(responseJSON(200, URL))

# https://address/md/api/mp4?id=hoge
@app.route('/md/api/mp4', methods=['POST'])
def downloadMp4():
	options = {
			'format': 'bestvideo[ext=mp4]/bestvideo',
			'outtmpl': './Music/%(title)s.%(ext)s'
			}

	ID = request.args.get("v")
	if ID == None:
		return jsonify(responseJSON(0, "None"))
	else:
		URL = "https://www.youtube.com/watch?v="+str(ID)
		downloadMedia(URL, options)
		return jsonify(responseJSON(200, URL))

if __name__ == "__main__":
	app.run(debug=True, host='0.0.0.0', port=5288)