# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# 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 argparse

import yaml
from api_base import PREFIX_TENSOR_NAME
from api_gen import (
    ForwardAPI,
    api_namespace,
    declare_extension_api,
    header_include,
    source_include,
)

######################
# Code Gen Templates #
######################

API_IMPL_TEMPLATE = """
PADDLE_API {} {}({}) {{
  // Kernel Key Construction{}
  // Kernel Dispatch Body{}
}}
"""
DIPATCH_END_GUARD_TEMPLATE = """
PADDLE_THROW(phi::errors::Unimplemented(
          "The kernel of ({}) for input tensors is unimplemented, please check the type of input tensors."));
"""

# TODO(chenweihang): add profile function code later
# TODO(chenweihang): add view support later
MAIN_DIST_BRANCH_TEMPLATE = """
  // Auto Parallel condition
  if ({}) {{
    // 1. Create API Output & Prepare Dist and Dense Output{}
    // 2. InferSPMD (Infer Global Shape and DistAttr of Inputs&Outputs){}
    // 3. Select Kernel{}
    // 4. Reshard Input{}
    // 5. PrepareData (DataTransform & Prepare Dist and Dense Input){}
    // 6. Infer Local DenseTensor Meta{}
    // 7. DenseTensor Kernel Call{}
    // 8. Reshard Output{}
    // 9. Return
    {}
  }}
"""

# Auto Parallel condition
AUTO_PARALLEL_COND_TEMPLATE = """AllInputsAreDistTensor({})"""

# 1. Create API Outputs
API_OUT_CREATION_TEMPLATE = """
    {} api_output{};
"""
INPLACE_API_OUT_CREATION_TEMPLATE = """
    {} api_output{{{}}};
"""
SINGLE_OUT_CREATION_TEMPLATE = """
    auto dist_out = SetKernelDistOutput(&api_output);
    auto dense_out = dist_out->mutable_value();
"""
MULTI_SINGLE_OUT_CREATION_TEMPLATE = """
    auto dist_out_{} = SetKernelDistOutput({});
    auto dense_out_{} = dist_out_{}->mutable_value();
"""

# TODO(chenweihang): support vector and tuple output later
VECTOR_OUT_CREATION_TEMPLATE = """
"""
MULTI_VECTOR_OUT_CREATION_TEMPLATE = """
    auto dist_out_{} = {}({}, {});
    auto dense_out_{} = dist_out_{}->mutable_value();
"""
TUPLE_OUT_CREATION_TEMPLATE = """
"""

# 2. InferSPMD
# Call InferMeta now, replace by InferSPMD function later
# TODO(chenweihang): InferSPMD function design
SINGLE_DIST_META_IN_TEMPLATE = """MakeMetaTensor(*{}.impl()), """
# TODO(chenweihang): support vector and optional args later
VECTOR_DIST_META_IN_TEMPLATE = """
"""
OPTIONAL_DIST_VECTOR_META_IN_TEMPLATE = """
"""
SINGLE_DIST_META_OUT_DECL_TEMPLATE = """
    phi::MetaTensor meta_{}({});"""
INFER_SPMD_TEMPLATE = """
    phi::{}({}{});
"""

# 3. Select Kernel
KERNEL_SELECTION_TEMPLATE = """
    VLOG(6) << "{} API dist branch: kernel key: [" << kernel_backend << ", " << kernel_layout << ", "<< kernel_data_type << "]";
    auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
        "{}", {{kernel_backend, kernel_layout, kernel_data_type}});
    const auto& kernel = kernel_result.kernel;
    VLOG(6) << "{} kernel: " << kernel;
    auto* dev_ctx = GetDeviceContextByBackend(kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
"""

# 4. Reshard Input
INPUT_RESHARD_TEMPLATE = """
"""

# 5. PrepareData
SINGLE_PREPARE_DATA_TEMPLATE = """
    auto dist_input_{} = PrepareDataForDistTensor({}, GetKernelInputArgDef(kernel.InputAt({}), kernel_backend), {}, kernel_result.is_stride_kernel);
    auto input_{} = dist_input_{}->mutable_value();
"""
INFER_META_SINGLE_INPUT_TEMPLATE = """
    auto dist_input_{} = {}.impl();
    auto input_{} = static_cast<phi::distributed::DistTensor*>(dist_input_{}.get())->mutable_value();
"""
INFER_META_OPTIONAL_INPUT_TEMPLATE = """
    paddle::optional<phi::TensorBase> input_{} = {} ? paddle::optional<phi::TensorBase>(*{}->impl()) : paddle::none;
"""
INFER_META_VECTOR_INPUT_TEMPLATE = """
    auto input_{}_uq_ptr = TensorToDenseTensor({});
    const auto& input_{} = *input_{}_uq_ptr;
"""

# 6. Infer Local DenseTensor Meta
SINGLE_META_IN_TEMPLATE = """MakeMetaTensor(*input_{}), """
# TODO(chenweihang): support vector and optional args later
VECTOR_META_IN_TEMPLATE = """
"""
OPTIONAL_VECTOR_META_IN_TEMPLATE = """
"""
SINGLE_META_OUT_DECL_TEMPLATE = """
    phi::MetaTensor meta_{}({});"""
INFER_META_TEMPLATE = """
    phi::{}({}{});
"""

# 7. DenseTensor Kernel Call
# TODO(chenweihang): support kernel fallback later
SINGLE_OUTPUT_NAME = """dense_out"""
# TODO(chenweihang): support vector and tuple output later
VECTOR_OUTPUT_NAME_TEMPLATE = """
"""
TUPLE_OUTPUT_NAME_TEMPLATE = """
"""
KERNEL_CALL_TEMPLATE = """
    using kernel_signature = {};
    auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
    (*kernel_fn)({}, {});
"""

# 8. Reshard Output
OUTPUT_RESHARD_TEMPLATE = """
"""

# BaseAPI members:
# inputs:
#     names : [], list of input names
#     input_info : {input_name : type}
# attrs:
#     names : [], list of attribute names
#     attr_info : { attr_name : (type, default_values)}
# outputs:
#     names : [], list of output names
#     types : [], list of output types
#     out_size_expr : [], expression for getting size of vector<Tensor>

class DistForwardAPI(ForwardAPI):
    def __init__(self, api_item_yaml):
        super().__init__(api_item_yaml)
        self.init_dist_api_members()

    def init_dist_api_members(self):
        self.gene_dist_input_func = {
            "const Tensor&": {
                "dense": self.generate_dense_input,
            },
            "const paddle::optional<Tensor>&": {
                "dense": self.generate_dense_input,
            },
        }

        self.inplace_flag = False
        self.dist_output_args = []
        self.dense_output_args = []

    def need_to_generate_code_for_inplace_impl(self, i):
        return (
            self.inplace_flag
            and self.inplace_map is not None
            and self.outputs['names'][i] in self.inplace_map
        )

    def need_to_generate_code_for_view_impl(self, i):
        return (
            not self.inplace_flag
            and self.view_map is not None
            and self.outputs['names'][i] in self.view_map
        )

    def is_inplace_output(self, i):
        return self.outputs['names'][i] in self.inplace_map

    def is_inplace_and_optional_output(self, i):
        return (
            self.outputs['names'][i] in self.inplace_map
            and self.inplace_map[self.outputs['names'][i]] in self.optional_vars
        )

    def vector_output_size_assertion_check(self):
        assert (
            self.outputs['out_size_expr'] is not None
        ), f"{self.api}: The out size expr : '{{expr}}' should be set when output has Tensor[]. You can refer 'split' api."

    def generate_if_condition_code(self) -> str:
        input_args = ""
        for input_name in self.inputs['names']:
            input_args = input_args + input_name + ", "
        if len(input_args) > 2:
            input_args = input_args[:-2]
        return AUTO_PARALLEL_COND_TEMPLATE.format(input_args)

    def generate_output_creation_code(self) -> str:
        # forward api need to generate api and kernel outputs
        output_num = len(self.outputs['types'])
        return_type = self.get_return_type_with_intermediate(self.inplace_flag)
        output_creation_code = ""
        if output_num == 1:
            # api output generate
            if self.need_to_generate_code_for_inplace_impl(0):
                inplace_assign_code = (
                    " = " + self.inplace_map[self.outputs['names'][0]]
                )
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, inplace_assign_code
                )
            else:
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, ""
                )
            # kernel output generate
            self.dist_output_args.append('dist_out')
            self.dense_output_args.append('dense_out')
            if self.outputs['types'][0] == 'Tensor':
                output_creation_code += SINGLE_OUT_CREATION_TEMPLATE
            else:
                self.vector_output_size_assertion_check()
        elif output_num > 1:
            # api output generate
            if self.inplace_flag:
                inplace_assign_code = ""
                for i, out_name in enumerate(self.outputs['names']):
                    if self.need_to_generate_code_for_inplace_impl(i):
                        inplace_assign_code += self.inplace_map[out_name] + ', '
                    else:
                        inplace_assign_code += 'Tensor(), '
                inplace_assign_code = inplace_assign_code[:-2]
                output_creation_code += (
                    INPLACE_API_OUT_CREATION_TEMPLATE.format(
                        return_type, inplace_assign_code
                    )
                )
            else:
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, ""
                )
            # kernel output generate
            for i, out_type in enumerate(self.outputs['types']):
                self.dist_output_args.append(f'dist_out_{i}')
                self.dense_output_args.append(f'dense_out_{i}')
                set_out_func = "SetKernelDistOutput"
                get_out_code = f"&std::get<{i}>(api_output)"
                if self.is_inplace_and_optional_output(i):
                    get_out_code = f"std::get<{i}>(api_output).get_ptr()"

                if out_type == 'std::vector<Tensor>':
                    self.vector_output_size_assertion_check(i)
                    # Special case for inplace vector and inplace optional<vector>
                    # TODO(chenweihang): support this branch later
                    if self.is_inplace_output():
                        set_out_func = "SetInplaceVectorKernelOutput"
                        if self.is_inplace_and_optional_output(i):
                            set_out_func = (
                                "SetInplaceOptionalVectorKernelOutput"
                            )
                            get_out_code = f"std::get<{i}>(api_output)"
                    output_creation_code += (
                        MULTI_VECTOR_OUT_CREATION_TEMPLATE.format(
                            i,
                            set_out_func,
                            self.outputs['out_size_expr'][i],
                            get_out_code,
                            i,
                            i,
                        )
                    )
                else:
                    output_creation_code += (
                        MULTI_SINGLE_OUT_CREATION_TEMPLATE.format(
                            i, get_out_code, i, i
                        )
                    )
        else:
            raise ValueError(
                "{} : Output error: the output should not be empty.".format(
                    self.api
                )
            )

        return output_creation_code

    def generate_infer_spmd_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        output_names = self.outputs['names']

        # 1. get infer meta func name
        infer_meta = self.infer_meta
        infer_meta_func_code = infer_meta['func']

        # 2. get meta tensor input args
        infer_meta_params = (
            infer_meta['param']
            if infer_meta['param'] is not None
            else input_names + attr_names
        )
        input_args_code = ""
        for param in infer_meta_params:
            if param in input_names:
                if self.inputs['input_info'][param] == "const Tensor&":
                    input_args_code += SINGLE_DIST_META_IN_TEMPLATE.format(
                        param
                    )
                else:
                    raise ValueError(
                        f"{self.api} : Param of infer_spmd error : {self.inputs['input_info'][param]} type is not supported."
                    )
            elif param in attr_names:
                input_args_code = input_args_code + param + ", "
            elif isinstance(param, str):
                input_args_code = input_args_code + "\"" + param + "\", "
            elif isinstance(param, bool):
                input_args_code = input_args_code + str(param).lower() + ", "
            else:
                input_args_code = input_args_code + str(param) + ", "

        # 3. get meta tensor output args
        output_decl_code = ""
        output_args_code = ""
        for i, out_name in enumerate(self.dist_output_args):
            if self.outputs['types'][i] == 'std::vector<Tensor>':
                # TODO(chenweihang): support vector output later
                pass
            else:
                output_decl_code += SINGLE_DIST_META_OUT_DECL_TEMPLATE.format(
                    out_name, out_name
                )
                if len(self.dense_output_args) == 1:
                    output_args_code += f"&meta_{out_name}, "
                else:
                    output_args_code += (
                        f"{out_name} ? &meta_{out_name} : nullptr, "
                    )
        output_args_code = output_args_code[:-2]

        return output_decl_code + INFER_SPMD_TEMPLATE.format(
            infer_meta_func_code, input_args_code, output_args_code
        )

    def generate_kernel_selection_code(self) -> str:
        return KERNEL_SELECTION_TEMPLATE.format(
            self.api, self.kernel['func'][0], self.kernel['func'][0]
        )

    def generate_reshard_input_code(self) -> str:
        return INPUT_RESHARD_TEMPLATE.format()

    # override BaseAPI's method
    def generate_dense_input(
        self,
        input_name,
    ):
        input_tensor_code = ""
        trans_flag = self.gene_trans_flag(input_name)
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        kernel_param = self.kernel['param']
        if kernel_param is None:
            kernel_param = input_names + attr_names

        input_tensor_code += SINGLE_PREPARE_DATA_TEMPLATE.format(
            input_name,
            input_name,
            kernel_param.index(input_name),
            trans_flag,
            input_name,
            input_name,
        )

        return input_tensor_code

    def generate_prepare_data_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        kernel_param = self.kernel['param']
        if kernel_param is None:
            kernel_param = input_names + attr_names
        input_tensor_code = ""
        for i, input_name in enumerate(input_names):
            # set input code
            if input_name in kernel_param:
                # onlu support dense tensor
                api_tensor_type = self.inputs['input_info'][input_name]
                phi_tensor_type = 'dense'
                if api_tensor_type in self.gene_dist_input_func.keys():
                    input_tensor_code += self.gene_dist_input_func[
                        api_tensor_type
                    ][phi_tensor_type](input_name)
                else:
                    # do nothing
                    pass
            else:
                if input_name in self.infer_meta['param']:
                    if input_name in self.optional_vars:
                        input_tensor_code += (
                            INFER_META_OPTIONAL_INPUT_TEMPLATE.format(
                                input_name, input_name, input_name, input_name
                            )
                        )
                    else:
                        if (
                            self.inputs['input_info'][input_name]
                            == "const std::vector<Tensor>&"
                        ):
                            input_tensor_code += (
                                INFER_META_VECTOR_INPUT_TEMPLATE.format(
                                    input_name, input_name, input_name
                                )
                            )
                        else:
                            input_tensor_code += (
                                INFER_META_SINGLE_INPUT_TEMPLATE.format(
                                    input_name,
                                    input_name,
                                    input_name,
                                    input_name,
                                )
                            )

        return input_tensor_code

    def generate_infer_meta_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        output_names = self.outputs['names']

        # 1. get infer meta func name
        infer_meta = self.infer_meta
        infer_meta_func_code = infer_meta['func']

        # 2. get meta tensor input args
        infer_meta_params = (
            infer_meta['param']
            if infer_meta['param'] is not None
            else input_names + attr_names
        )
        input_args_code = ""
        for param in infer_meta_params:
            if param in input_names:
                if self.inputs['input_info'][param] == "const Tensor&":
                    input_args_code += SINGLE_META_IN_TEMPLATE.format(param)
                else:
                    raise ValueError(
                        f"{self.api} : Param of infer_meta error : {self.inputs['input_info'][param]} type is not supported."
                    )
            elif param in attr_names:
                input_args_code = input_args_code + param + ", "
            elif isinstance(param, str):
                input_args_code = input_args_code + "\"" + param + "\", "
            elif isinstance(param, bool):
                input_args_code = input_args_code + str(param).lower() + ", "
            else:
                input_args_code = input_args_code + str(param) + ", "

        # 3. get meta tensor output args
        output_decl_code = ""
        output_args_code = ""
        for i, out_name in enumerate(self.dense_output_args):
            if self.outputs['types'][i] == 'std::vector<Tensor>':
                # TODO(chenweihang): support vector output later
                pass
            else:
                output_decl_code += SINGLE_META_OUT_DECL_TEMPLATE.format(
                    out_name, out_name
                )
                if len(self.dense_output_args) == 1:
                    output_args_code += f"&meta_{out_name}, "
                else:
                    output_args_code += (
                        f"{out_name} ? &meta_{out_name} : nullptr, "
                    )
        output_args_code = output_args_code[:-2]

        return output_decl_code + INFER_META_TEMPLATE.format(
            infer_meta_func_code, input_args_code, output_args_code
        )

    def generate_kernel_call_code(self) -> str:
        dense_input_trans_map = {
            'const Tensor&': 'const phi::DenseTensor&',
            'const std::vector<Tensor>&': 'const std::vector<const phi::DenseTensor*>&',
            'const paddle::optional<Tensor&>': 'paddle::optional<const phi::DenseTensor&>',
            'const paddle::optional<Tensor>&': 'const paddle::optional<phi::DenseTensor>&',
            'const paddle::optional<std::vector<Tensor>>&': 'const paddle::optional<std::vector<const phi::DenseTensor*>>&',
        }
        dense_output_trans_map = {
            'Tensor': 'phi::DenseTensor*',
            'std::vector<Tensor>': 'std::vector<phi::DenseTensor*>',
        }

        input_names = self.inputs['names']
        input_infos = self.inputs['input_info']
        kernel_args_type_list = ['const phi::DeviceContext&']

        attr_names = self.attrs['names']
        kernel_args = self.kernel['param']
        if kernel_args is None:
            kernel_args = input_names + attr_names

        # 1. generate input args list
        input_args = ["*dev_ctx"]
        for arg in kernel_args:
            if arg in input_names:
                if arg in self.optional_vars:
                    input_args.append(PREFIX_TENSOR_NAME + arg)
                else:
                    if input_infos[arg] == "const Tensor&":
                        input_args.append("*" + PREFIX_TENSOR_NAME + arg)
                    elif input_infos[arg] == "const std::vector<Tensor>&":
                        input_args.append(PREFIX_TENSOR_NAME + arg)
                    else:
                        # do nothing
                        pass
                kernel_args_type_list.append(
                    dense_input_trans_map[input_infos[arg]]
                )
            elif arg in attr_names:
                if 'IntArray' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append('const phi::IntArray&')
                    arg = 'phi::IntArray(' + arg + ')'
                elif 'vector<phi::Scalar>' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append(
                        'const std::vector<phi::Scalar>&'
                    )
                elif 'Scalar' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append('const phi::Scalar&')
                    arg = 'phi::Scalar(' + arg + ')'
                else:
                    kernel_args_type_list.append(
                        self.attrs['attr_info'][arg][0]
                    )
                input_args.append(arg)
            elif isinstance(arg, bool):
                input_args.append(str(arg).lower())
            else:
                input_args.append(str(arg))

        # 2. generate output args list
        # record into `self.dense_output_args` in `generate_output_creation_code` function

        # 3. generate kernel signature
        for i, out_type in enumerate(self.outputs['types']):
            kernel_args_type_list.append(dense_output_trans_map[out_type])
        kernel_signature = "void(*)(" + ", ".join(kernel_args_type_list) + ")"

        return KERNEL_CALL_TEMPLATE.format(
            kernel_signature,
            ", ".join(input_args),
            ", ".join(self.dense_output_args),
        )

    def generate_reshard_output_code(self) -> str:
        return OUTPUT_RESHARD_TEMPLATE.format()

    def generate_return_code(self) -> str:
        return self.gene_return_code()

    def generate_auto_paralel_branch(self) -> str:
        # if no tensor input, do not genetate auto parallel branch
        if len(self.inputs['names']) == 0:
            return ""
        return MAIN_DIST_BRANCH_TEMPLATE.format(
            self.generate_if_condition_code(),
            self.generate_output_creation_code(),
            self.generate_infer_spmd_code(),
            self.generate_kernel_selection_code(),
            self.generate_reshard_input_code(),
            self.generate_prepare_data_code(),
            self.generate_infer_meta_code(),
            self.generate_kernel_call_code(),
            self.generate_reshard_output_code(),
            self.generate_return_code(),
        )

    def check_argument_whether_support_auto_parallel(self):
        for name in self.inputs['names']:
            if self.inputs['input_info'][name] != "const Tensor&":
                return False
        for out_type in self.outputs['types']:
            if out_type != "Tensor":
                return False
        return True

    # override BaseAPI's method
    def gene_base_api_code(self, inplace_flag=False):
        # init status
        self.inplace_flag = inplace_flag
        self.dist_output_args = []
        self.dense_output_args = []
        # generate api body
        api_func_name = self.get_api_func_name()
        if inplace_flag and api_func_name[-1] != '_':
            api_func_name += '_'

        if len(self.kernel['func']) > 1:
            kernel_dispatch_code = ''
            for kernel_name in self.kernel['func']:
                kernel_dispatch_code += self.gene_dispatch_code(
                    kernel_name, inplace_flag
                )
            return API_IMPL_TEMPLATE.format(
                self.get_return_type(inplace_flag),
                api_func_name,
                self.get_define_args(inplace_flag),
                self.gene_kernel_select(),
                kernel_dispatch_code
                + DIPATCH_END_GUARD_TEMPLATE.format(self.api),
            )
        else:
            # auto parallel branch, all apis contains this branch default
            # 1. only works for the ops contains single kernel
            # 2. doesn't support initialize ops now
            # 3. doesn't support view api
            # 4. only for general forward and backward
            # 5. only support single tensor input and output
            dist_branch_code = ""
            if (
                len(self.inputs['names']) > 0
                and len(self.view_map) == 0
                and self.check_argument_whether_support_auto_parallel()
            ):
                dist_branch_code = self.generate_auto_paralel_branch()

            return API_IMPL_TEMPLATE.format(
                self.get_return_type(inplace_flag),
                api_func_name,
                self.get_define_args(inplace_flag),
                self.gene_kernel_select(),
                dist_branch_code
                + self.gen_kernel_code(
                    self.kernel['func'][0], '', inplace_flag
                ),
            )

def generate_api(
    api_yaml_path, is_fused_ops_yaml, header_file_path, source_file_path
):
    apis = []

    for each_api_yaml in api_yaml_path:
        with open(each_api_yaml, 'r') as f:
            api_list = yaml.load(f, Loader=yaml.FullLoader)
            if api_list:
                apis.extend(api_list)

    header_file = open(header_file_path, 'w')
    source_file = open(source_file_path, 'w')

    namespace = api_namespace()

    header_file.write("#pragma once\n")
    header_file.write(header_include())
    header_file.write(namespace[0])

    include_header_file = (
        "paddle/phi/api/include/fused_api.h"
        if is_fused_ops_yaml is True
        else "paddle/phi/api/include/api.h"
    )
    # not all fused ops supoort dygraph
    if is_fused_ops_yaml is True:
        new_apis = [
            api
            for api in apis
            if "support_dygraph_mode" in api
            and api["support_dygraph_mode"] is True
        ]
        apis = new_apis

    source_file.write(source_include(include_header_file))
    source_file.write(namespace[0])

    for api in apis:
        dist_foward_api = DistForwardAPI(api)
        if dist_foward_api.is_dygraph_api:
            dist_foward_api.is_dygraph_api = False

        header_file.write(dist_foward_api.gene_api_declaration())
        if is_fused_ops_yaml is True:
            source_file.write(dist_foward_api.gene_api_code())
        else:
            source_file.write(dist_foward_api.gene_api_code())

    header_file.write(namespace[1])
    source_file.write(namespace[1])

    source_file.write(declare_extension_api())

    header_file.close()
    source_file.close()

def main():
    parser = argparse.ArgumentParser(
        description='Generate PaddlePaddle C++ API files'
    )
    parser.add_argument(
        '--api_yaml_path',
        help='path to api yaml file',
        nargs='+',
        default=['paddle/phi/api/yaml/ops.yaml'],
    )

    parser.add_argument(
        '--is_fused_ops_yaml',
        help='flag of fused ops yaml',
        action='store_true',
    )

    parser.add_argument(
        '--api_header_path',
        help='output of generated api header code file',
        default='paddle/phi/api/include/api.h',
    )

    parser.add_argument(
        '--api_source_path',
        help='output of generated api source code file',
        default='paddle/phi/api/lib/api.cc',
    )

    options = parser.parse_args()

    api_yaml_path = options.api_yaml_path
    is_fused_ops_yaml = options.is_fused_ops_yaml
    header_file_path = options.api_header_path
    source_file_path = options.api_source_path

    generate_api(
        api_yaml_path, is_fused_ops_yaml, header_file_path, source_file_path
    )

if __name__ == '__main__':
    main()


import z3
import ast
import logging
#from pyState import z3Helpers
from pyObjectManager.Int import Int
from pyObjectManager.Real import Real
from pyObjectManager.BitVec import BitVec
from pyObjectManager.List import List
from pyObjectManager.Ctx import Ctx
from pyObjectManager.String import String
from pyObjectManager.Char import Char
import pyState

logger = logging.getLogger("ObjectManager")

CTX_GLOBAL = 0
CTX_RETURNS = 1

class ObjectManager:
    """
    Object Manager will keep track of objects. Generally, Objects will be variables such as ints, lists, strings, etc.
    """

    def __init__(self,variables=None,returnObjects=None,state=None):
        self.variables = {CTX_GLOBAL: Ctx(CTX_GLOBAL), CTX_RETURNS: Ctx(CTX_RETURNS)} if variables is None else variables
        self.returnObjects = returnObjects if returnObjects is not None else {}

        if state is not None:
            self.setState(state)

    def setState(self,state):
        """
        This is a bit strange, but State won't copy correctly due to Z3, so I need to bypass this a bit by setting State each time I copy
        """
        assert type(state) == pyState.State

        self.state = state
        for ctx in self.variables:
            self.variables[ctx].setState(state)

    def newCtx(self,ctx):
        """
        Sets up a new context (ctx)
        """
        assert ctx is not None

        self.variables[ctx] = Ctx(ctx)
        self.variables[ctx].setState(self.state)

    def setVar(self,varName,ctx,var):
        """
        Input:
            varName = variable name (i.e.: 'x')
            ctx = Context to set for
            var = variable object of type pyObjectManager.X
        Action:
            Sets variable to the input (var) object
        Returns:
            Nothing
        """
        assert type(varName) is str
        assert type(ctx) is int
        assert type(var) in [Int, Real, BitVec, List]

        self.variables[ctx][varName] = var

    def getVar(self,varName,ctx,varType=None,kwargs=None,softFail=None):
        """
        Input:
            varName = name of variable to get
            ctx = Context for variable
            (optional) varType = Class type of variable (ex: pyObjectManager.Int)
            (optional) kwargs = args needed to instantiate variable
            (optional) softFail = True/False, should raise an exception if getVar fails. Default is False
        Action:
            Find appropriate variable object, creating one if necessary
        Returns:
            pyObjectManager object for given variable (i.e.: pyObjectManager.Int)
        """

        # Attempt to return variable
        assert type(varName) is str
        assert type(ctx) is int
        assert varType in [None, Int, Real, BitVec, List, String, Char]

        softFail = False if softFail is None else softFail
        create = False
        count = None

        # Check that we already have this variable defined
        if varName in self.variables[ctx]:

            # Check the type of the var is correct
            if varType is not None:

                # If the variable type is different or it's settings are different, we need to create a new object
                if type(self.variables[ctx][varName]) is not varType or not self.variables[ctx][varName]._isSame(**kwargs if kwargs is not None else {}):
                    create = True
                    # Re-using variable names is BAD!
                    count = self.variables[ctx][varName].count + 1

            # If we can just return the current one, let's do it
            if not create:
                return self.variables[ctx][varName]

        # Looks like we need to create a var
        if varType == None:
            if not softFail:
                err = "getVar: Need to create '{0}' but no type information given".format(varName)
                logger.error(err)
                raise Exception(err)
            else:
                # We're soft failing, just return None
                return None

        # Make the var
        self.variables[ctx][varName] = varType(varName=varName,ctx=ctx,count=count,state=self.state,**kwargs if kwargs is not None else {})

        return self.variables[ctx][varName]

    def getParent(self,key,haystack=None):
        """
        Returns the parent object for any given object by recursively searching.
        """
        # TODO: This might get to be a long search if there are a lot of variables...

        haystack = self.variables if haystack is None else haystack

        if type(haystack) in [dict, Ctx]:
            for k,v in haystack.items():
                if v == key:
                    return haystack
                elif type(v) in [dict, List, Ctx]:
                    p = self.getParent(key,v)
                    if p:
                        return p
        elif isinstance(haystack,List):
            for v in haystack:
                if v == key:
                    return haystack
                elif type(v) in [dict,List]:
                    p = self.getParent(key,v)
                    if p:
                        return p

    def copy(self):
        """
        Return a copy of the Object Manager
        """

        return ObjectManager(
            variables = {key:self.variables[key].copy() for key in self.variables},
            returnObjects = {key:self.returnObjects[key].copy() for key in self.returnObjects}
        )


import argparse
import click
import pandas as pd
import re
from datetime import datetime
from datetime import timedelta
from os import chdir
from os import chdir
from os import getcwd
from os import makedirs
from os import mkdir
from os import remove
from os import rename
from os import system
from os.path import basename
from os.path import dirname
from os.path import expanduser
from os.path import getmtime
from os.path import isdir
from os.path import isdir
from os.path import isfile
from os.path import join
from os.path import splitext
from pydoni.db import progrun_update
from pydoni.os import listfiles
from pydoni.os import macos_notify
from pydoni.sh import syscmd
from pydoni.vb import echo
from pydoni.vb import macos_notify
from pydoni.vb import program_complete
from pydoni.vb import verbose_header
from send2trash import send2trash
from shutil import copy2
from tqdm import tqdm

def parse_date(date):
    """
    Get date in format YYYY-MM-DD if entered as 'today' or 'yesterday'.

    Arguments:
        date {str} -- date in format YYYYMMDD or 'today' or 'yesterday'

    Returns:
        {str}
    """
    if date == 'today':
        date = datetime.strftime(datetime.now(), '%Y-%m-%d')
    elif date == 'yesterday':
        date = datetime.strftime(datetime.now() - timedelta(days=1), '%Y-%m-%d')
    assert re.search(r'^\d{4}-\d{2}-\d{2}$', date)
    return date

def parse_target_subcollection(fname):
    """
    Given a filename, get the target subcollection as one of ['Photo', 'Video', 'Gopro', 'Drone'].

    Arguments:
        fname {str} -- name of file to parse

    Returns:
        {str}
    """
    c = Convention()
    e = Extension()
    media_type = 'Photo' if splitext(fname)[1].lower() in e.photo else 'Video'
    conv = c.photo if media_type == 'Photo' else c.video
    camera = re.match(conv, fname).group('camera')
    if re.search(r'gopro|hero', camera, flags=re.IGNORECASE):
        return 'Gopro'
    elif 'FC' in camera or 'L1D' in camera:
        if media_type == 'Photo':
            return join('Drone', 'Photo')
        elif media_type == 'Video':
            return join('Drone', 'Video')
    else:
        return media_type

class Extension(object):
    """
    Hold valid file extension lists of strings.
    """

    def __init__(self):
        self.photo = ['.jpg', '.jpeg', '.dng', '.arw', '.cr2']
        self.video = ['.mov', '.mp4', '.mts', '.m4v', '.avi']

class Convention(object):
    """
    Hold naming convention regex strings.
    """

    def __init__(self):
        self.photo = ''
        self.video = ''
        conv_base = '' +\
            r'(?P<year>\d{4})' +\
            r'(?P<month>\d{2})' +\
            r'(?P<day>\d{2})' +\
            r'(?P<initials>[A-Za-z]{2})' +\
            r'(?P<hours>\d{2})' +\
            r'(?P<minutes>\d{2})' +\
            r'(?P<seconds>\d{2})'
        self.photo = conv_base +\
            '_' +\
            r'(?P<camera>.*?)' +\
            '_' +\
            r'(?P<seqnum>(\d{4,}|\d+-\d|-\d|))' +\
            r'(?P<affix>(-HDR-*\d*|-Pano-*\d*|-Edit-*\d*|-Stack-*\d*|)*)' +\
            r'(?P<ext>\.[a-z0-9]{3})'
        self.video = conv_base +\
            '_' +\
            r'(?P<camera>.*?)' +\
            '_' +\
            r'(?P<seqnum>.*(\d+))' +\
            '_' +\
            r'(?P<quality>Q\d+(K|P))' +\
            r'(?P<fps>\d{2,3}FPS)' +\
            r'(?P<ext>\.[a-z0-9]{3})'

class MediaFile(Convention, Extension):
    """
    Gather information for a photo or video file.

    Arguments:
        fpath {str} -- absolute path to file to initialize MediaFile class on

    Keyword Arguments:
        run_exiftool {bool} -- if True, run ExifTool automatically and save output dictionary as `self.exif` attribute (default: {False})
    """

    def __init__(self, fpath, run_exiftool=False):
        from os.path import isfile, dirname, basename, splitext, abspath
        from pydoni.classes import Attribute

        assert isinstance(fpath, str)
        assert isfile(fpath)
        assert isinstance(run_exiftool, bool)

        # Inherit naming conventions and extensions
        Convention.__init__(self)
        Extension.__init__(self)

        # Assign filename and directory name attributes to `self`
        self.fpath_abs = abspath(fpath)
        self.fname = basename(self.fpath_abs)
        self.dname = dirname(self.fpath_abs)
        assert self.fname != self.fpath_abs
        assert self.dname > ''

        # Define extensions
        self.ext = splitext(self.fname)[1].lower()
        self.valid_ext = Extension()
        assert self.ext in self.valid_ext.photo + self.valid_ext.video

        # Parse media type from extension
        self.mtype = self.parse_media_type()
        self.remove_flag = True if self.mtype == 'remove' else False

        # Run ExifTool if specified
        if run_exiftool is True:
            from pydoni.sh import EXIF
            self.exif = EXIF(self.fpath_abs).run(wrapper='doni')

    def parse_media_type(self):
        """
        Given a file extension, get the type of media (one of 'photo' or 'video')

        Returns:
            {str} -- type of media
        """

        if self.ext in self.valid_ext.photo:
            return 'photo'
        elif self.ext in self.valid_ext.video:
            return 'video'
        elif self.ext in self.valid_ext.rm:
            return 'remove'
        else:
            from pydoni.vb import echo
            echo('Invalid file extension slipped through the cracks!', fn_name='MediaFile.parse_media_type', abort=True)

    def build_fname(self, initials, tz_adjust=0):
        """
        Build new filename from EXIF metadata.

        Arguments:
            initials {str} -- two character initials string

        Keyword Arguments:
            tz_adjust {int} -- alter 'hours' by set number to account for timezone adjust (default: {0})

        Returns:
            {str} -- new filename according to convention (basename)
        """

        from os.path import join

        assert isinstance(initials, str)
        assert len(initials) == 2
        req_attrs = ['exif', 'mtype', 'dname', 'fname']
        for ra in req_attrs:
            assert hasattr(self, ra)
        assert self.mtype in ['photo', 'video']
        assert isinstance(tz_adjust, int)
        assert tz_adjust in range(-24, 24)

        def search_exif(exif, attrs, def_str):
            """
            Search exif dictionary for name attributes in the order specified. For example
            if `attrs` argument is ['image_width', 'video_size'], first search `exif` for
            the attribute 'image_width'. If it exists and is a valid value, return this
            value. If not, search for 'video_size' and return that value if it exists and
            is valid. If neither value exists, return an arbitrary string.

            Arguments:
                attrs {str} or {list} -- attribute name(s) to search exif dictionary for
                def_str {str} -- return this string if attributes not found
            """
            assert isinstance(exif, dict)
            assert isinstance(attrs, str) or isinstance(attrs, list)
            assert isinstance(def_str, str)

            attrs = [attrs] if isinstance(attrs, str) else attrs
            bogus = ['0000:00:00 00:00:00']

            for attr in attrs:
                if attr in exif.keys():
                    val = exif[attr]
                    if val not in bogus:
                        return val

            return def_str

        def extract_video_quality(exif, def_str):
            """
            Given EXIF dictionary, extract video quality.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- video quality string
            """
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)

            image_width = search_exif(
                exif    = exif,
                attrs   = ['image_width', 'video_size'],
                def_str = def_str)
            if image_width == def_str:
                return def_str
            elif image_width == '3840':
                return '4K'
            elif image_width == '2704':
                return '1520P'
            elif image_width == '1920':
                return '1080P'
            elif image_width == '1280':
                return '720P'
            else:
                return image_width

        def extract_video_framerate(exif, def_str):
            """
            Given EXIF dictionary, extract video frame rate.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- video framerate string
            """
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)

            framerate = search_exif(
                exif    = exif,
                attrs=['video_avg_frame_rate', 'video_frame_rate'],
                def_str = def_str)
            if framerate == def_str:
                return def_str
            else:
                try:
                    return round(float(framerate))
                except:
                    return framerate

        def extract_capture_date_and_time(exif, def_str, mtype, tz_adjust=None):
            """
            Given EXIF dictionary, extract file capture date and time.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found
                mtype {str} -- `MediaFile.mtype` (media type string)

            Keyword Arguments:
                tz_adjust {int} -- alter 'hours' by set number to account for timezone adjust

            Returns:
                {tuple} -- (capture date string, capture time string)
            """
            import re
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)
            assert mtype in ['photo', 'video']

            if mtype == 'photo':
                attrs = ['create_date', 'file_modification_date/time']
            elif mtype == 'video':
                attrs = ['file_modification_date/time', 'create_date']

            # Get capture date
            raw_dt = search_exif(
                exif    = exif,
                attrs   = attrs,
                def_str = def_str
            )
            capture_date = re.sub(r'(.*?)(\s+)(.*)', r'\1', raw_dt)
            capture_date = capture_date.replace(':', '')

            # Get capture time
            capture_time = raw_dt.replace(':', '')
            capture_time = re.sub(r'(.*?)(\s+)(.*)', r'\3', capture_time)[0:6]

            # Adjust hours by timezone if specified. May also need to alter day if timezone
            # adjust makes the day spill over into the next/previous day
            if tz_adjust is not None:
                import datetime
                dt = datetime.datetime.strptime(capture_date + capture_time, '%Y%m%d%H%M%S')
                dt = dt + datetime.timedelta(hours=tz_adjust)
                capture_date = dt.strftime('%Y%m%d')
                capture_time = dt.strftime('%H%M%S')

            return (capture_date, capture_time)

        def extract_sequence_number(fname, def_str):
            """
            Given filename, extract file sequence number as the final occurrence of 4 or 5
            digits in a filename.

            Arguments:
                fname {str} -- filename to extract sequence number from
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- sequence number
            """
            import re
            from os.path import basename
            assert isinstance(fname, str)
            assert isinstance(def_str, str)

            # Extract sequence number. Search first for 5 digits at the end of the filename.
            # If not found, search for 4 digits.

            # Clean up searchable area of filename by removing YYYYMMDDASHHMMSS and the string
            # for video files QXXXXFPS
            fname_ = basename(fname)
            fname_ = re.sub('_Q.*?FPS', '', fname_)
            fname_ = re.sub(r'^\d{8}\w{2}\d{6}_', '', fname_)

            # Search for 5 or 4 digits
            ptn = r'^(.*)(\d{5})(.*)$'
            if re.search(ptn, fname_):
                return re.sub(ptn, r'\2', fname_)
            else:
                ptn = r'^(.*)(\d{4})(.*)$'
                if re.search(ptn, fname_):
                    return re.sub(ptn, r'\2', fname_)
                else:
                    return def_str

        def extract_camera_model(exif, def_str, dname):
            """
            Given EXIF dictionary, extract camera model.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found
                dname {str} -- directory name of photo, only used in parsing camera model if all else fails

            Returns:
                {str} -- camera model string
            """
            from os.path import basename, isdir
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)
            assert isinstance(dname, str)
            assert isdir(dname)

            # Get camera model from raw exif
            cm_raw = search_exif(
                exif    = exif,
                attrs   = ['camera_model_name', 'device_model_name', 'model'],
                def_str = def_str
            )

            # If no camera model found, apply manual corrections
            if cm_raw == def_str:

                # Get make
                make_raw = search_exif(
                    exif    = exif,
                    attrs   = ['make', 'device_manufacturer', 'compressor_name'],
                    def_str = '{}'
                )
                valid_makes = ['Sony', 'Canon', 'DJI', 'Gopro', 'Gopro AVC Encoder', '{}']
                assert any([make_raw.lower() in [x.lower() for x in valid_makes]])
                for vm in valid_makes:
                    if vm.lower() in make_raw.lower():
                        make = vm

                # Manual corrections
                if make.lower() == 'dji' or basename(dname).lower() == 'drone':
                    return 'L1D-20c'  # DJI Mavic 2 Pro
                elif make.lower() == 'gopro' or basename(dname).lower() == 'gopro':
                    return 'HERO5 Black'
                else:
                    return def_str

            else:
                return cm_raw

        exif = self.exif
        fname = self.fname
        self.initials = initials

        # Default string to be returned by any of the extract* functions if that
        # attribute is unable to be found
        def_str = '{}'

        # Extract attributes from EXIF used regardless of photo or video
        self.cd, self.ct = extract_capture_date_and_time(exif, def_str, self.mtype, tz_adjust)
        self.sn = extract_sequence_number(fname, def_str)
        self.cm = extract_camera_model(exif, def_str, self.dname)

        # Build new filename
        if self.mtype == 'photo':
            newfname = "{}{}{}_{}_{}{}".format(
                self.cd, self.initials, self.ct, self.cm, self.sn, self.ext.lower())

        elif self.mtype == 'video':
            self.vq = extract_video_quality(exif, def_str)
            self.vfr = extract_video_framerate(exif, def_str)
            newfname = "{}{}{}_{}_{}_Q{}{}FPS{}".format(
                self.cd, self.initials, self.ct, self.cm, self.sn, self.vq, self.vfr, self.ext.lower())

        return newfname

    def convert_dng(self, remove_original=False):
        """
        Convert a photo file to dng.

        Keyword Arguments:
            remove_original {bool} -- if True, remove original raw file (default: {True})
        """
        from pydoni.sh import adobe_dng_converter
        assert self.ext in self.valid_ext.photo
        adobe_dng_converter(self.fpath_abs)
        if remove_original:
            from send2trash import send2trash
            send2trash(self.fpath_abs)

class DumpSD(object):
    """
    Dump photo and video files from a source volume to a target directory. Rename and Convert
    RAW files to DNG as necessary.

    Arguments:
        vol_source {str} -- source volume name
        coll_target {str} -- absolute directory path of destination collection
        date {str} -- dump photos with this creation date, format YYYY-MM-DD
        rename {bool} -- rename media files according to defined convention
        convert {bool} -- convert RAW photos to DNG
        verbose {bool} -- print messages to STDOUT
        notify {bool} -- send macOS notification upon program completion
        test_mode {bool} -- run program in test mode
        rc_initials {str} -- initials parameter for RenameConvert. Only must be specified if 'rename' or 'convert' is True (default: {AS})
        rc_tz_adjust {int} -- tz_adjust parameter for RenameConvert. Only must be specified if 'rename' or 'convert' is True (default: {0})
    """

    def __init__(
        self,
        vol_source,
        coll_target,
        date,
        rename,
        convert,
        verbose,
        notify,
        test_mode,
        rc_initials,
        rc_tz_adjust
        ):
        self.vol_source = join('/', 'Volumes', vol_source)
        self.coll_target = coll_target
        self.vol_target = '/'.join(self.coll_target.split('/')[0:3])
        self.date = parse_date(date)
        self.rename = rename
        self.convert = convert
        self.verbose = verbose
        self.notify = notify
        self.test_mode = test_mode
        self.rc_initials = rc_initials
        self.rc_tz_adjust = rc_tz_adjust

        assert isdir(self.vol_source)
        assert isdir(self.coll_target)
        assert isdir(self.vol_target)

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):

        prog_start = datetime.now()
        res = {'source': [], 'target': []}

        if self.test_mode:
            echo("Initializing Program: 'Dump SD' in Test Mode...", sleep=1)
            echo('')
            echo('Begin Operational Code')
            echo('')

        if self.test_mode:
            echo('Checking source and target directories...', sleep=.25, fg='blue')
            echo('Source volume name     : {} {}'.format(
                self.vol_source, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('Source volume type     : {} {}'.format(
                self.vol_source_type, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('Target collection      : {} {}'.format(
                self.coll_target, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('')

        # List all files in source volume that match given date
        chdir(self.vol_source)
        ext = Extension()
        fnames = [f for f in listfiles(recursive=True, ext=ext.photo+ext.video) \
            if datetime.fromtimestamp(getmtime(f)).strftime('%Y-%m-%d') == self.date \
            and not re.search(r'T\d{2,3}\.', f)  # Eliminate Sony movie thumbnail files
        ]
        if not len(fnames):
            echo("No photo or video files found at '{}' for day '{}'!".format(
                self.vol_source, self.date), abort=True)

        # Copy files to temporary directory
        tmpdir = join(self.vol_target, '.tmp.dumpsd')
        if isdir(tmpdir):
            send2trash(tmpdir)
        mkdir(tmpdir)
        if self.verbose:
            verbose_header('Dumping %s media files' % str(len(fnames)))
            with tqdm(total=len(fnames), unit='mediafile') as pbar:
                for f in fnames:
                    pbar.set_postfix(mediafile=f[-10:])
                    copy2(f, join(tmpdir, basename(f)))
                    res['source'].append(f)
                    pbar.update(1)
        else:
            for f in fnames:
                res['source'].append(f)
                copy2(f, join(tmpdir, basename(f)))

        # Run RenameConvert on directory
        if self.rename or self.convert:
            ignore_rename = not self.rename
            ignore_convert = not self.convert
            RenameConvert(
                dname=tmpdir,
                initials=self.rc_initials,
                tz_adjust=self.rc_tz_adjust,
                ignore_rename=ignore_rename,
                ignore_convert=ignore_convert,
                recursive=False,
                notify=self.notify,
                verbose=self.verbose
            ).run()

        # Copy files to target collection
        chdir(tmpdir)
        media_files = [x for x in listfiles() if splitext(x)[1].lower() != '.arw']
        dumped_subcolls = []
        for f in media_files:
            subcoll = parse_target_subcollection(f)
            if subcoll == basename(self.coll_target):
                destf = join(self.coll_target, f)
            else:
                destf = join(self.coll_target, subcoll, f)
            if not isdir(dirname(destf)):
                makedirs(dirname(destf))
            rename(f, destf)
            dumped_subcolls.append(subcoll)
            res['target'].append(f)

        if isdir(tmpdir):
            send2trash(tmpdir)

        if self.verbose:
            echo('')
            echo('Dump Summary', underline=True, fg='white', bold=True)
            dumped_subcolls = pd.Series(dumped_subcolls).value_counts()
            echo('Collection: %s' % self.coll_target, indent=2)
            for subcol, count in dumped_subcolls.iteritems():
                echo("Subcollection '{}': {} media files dumped".format(
                    click.style(subcol, fg='cyan'),
                    str(count)),
                    indent=2)

        if self.notify:
            content_image = '/Users/Andoni/GDrive/Code/git-doni/pydoni/pydoni/scripts/dump_sd/img/Dump-SD-Icon.png'
            content_image = content_image if isfile(content_image) else None
            macos_notify(title='Dump SD', message='Dump completed successfully!',
                content_image=content_image)

        progrun_update(
            name='Photo.dump_sd',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res=res)

class RenameConvert(object):
    """
    Rename/Convert program master class.

    Arguments:
        dname {str} -- directory name to run program on
        initials {str} -- two-character initial string
        tz_adjust {int} -- subtract this many hours from photo created date to use in file renaming (default: {0})
        ignore_rename {bool} -- if True, ignore the Rename portion of program (default: {False})
        ignore_convert {bool} -- if True, ignore the Convert portion of program (default: {False})
        recursive {bool} -- if True, operate on all subdirectories found in `dname` (default: {False})
        verbose {bool} -- if True, print messages to STDOUT (default: {False})
    """

    def __init__(
        self,
        dname,
        initials='AS',
        tz_adjust=0,
        ignore_rename=False,
        ignore_convert=False,
        recursive=False,
        verbose=False,
        notify=False
        ):
        self.dname = dname
        self.initials = initials
        self.tz_adjust = tz_adjust
        self.ignore_rename = ignore_rename
        self.ignore_convert = ignore_convert
        self.recursive = recursive
        self.verbose = verbose
        self.notify = notify

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):

        prog_start = datetime.now()
        res = {'rename': [], 'convert': []}

        assert isdir(self.dname)
        assert isinstance(self.initials, str)
        assert len(self.initials) == 2
        assert isinstance(self.tz_adjust, int)
        assert isinstance(self.ignore_rename, bool)
        assert isinstance(self.ignore_convert, bool)
        assert isinstance(self.recursive, bool)
        assert isinstance(self.verbose, bool)

        chdir(self.dname)
        c = Convention()

        # Rename
        if not self.ignore_rename:
            media_files_r = listfiles(full_names=True, recursive=self.recursive)
            media_files_r = [x for x in media_files_r \
                if not re.match(c.photo, basename(x)) \
                and not re.match(c.video, basename(x))]

            if len(media_files_r):

                if self.verbose:
                    print()
                    verbose_header('Renaming %s media files' % str(len(media_files_r)))
                    pbar = tqdm(total=len(media_files_r), unit='mediafile')

                for f in media_files_r:
                    if self.verbose:
                        pbar.set_postfix(mediafile=f[-10:])

                    mf = MediaFile(f, run_exiftool=True)

                    if mf.remove_flag is True:
                        send2trash(f)
                        if self.verbose:
                            pbar.update(1)
                        continue

                    newfname = mf.build_fname(initials=self.initials, tz_adjust=self.tz_adjust)
                    if basename(f) != basename(newfname):
                        rename(f, newfname)
                        res['rename'].append(f)

                    if self.verbose:
                        pbar.update(1)

            else:
                if self.verbose:
                    echo('No files to rename!', fg='green')

        # Convert to DNG
        if not self.ignore_convert:
            media_files_c = listfiles(full_names=True, ext=['arw', 'cr2'], recursive=self.recursive)

            if len(media_files_c):
                if self.verbose:
                    print()
                    verbose_header('Converting %s raw files' % str(len(media_files_c)))
                    pbar = tqdm(total=len(media_files_c), unit='rawfile')

                for f in media_files_c:
                    if self.verbose:
                        pbar.set_postfix(mediafile=f[-10:])

                    mf = MediaFile(f, run_exiftool=False).convert_dng(remove_original=True)
                    res['convert'].append(f)

                    if self.verbose:
                        pbar.update(1)

            else:
                if self.verbose:
                    echo('No files to convert!', fg='green')

        if self.verbose:
            echo('')
            echo('Rename/Convert Summary', underline=True, fg='white', bold=True)
            echo('Renamed media files: %s' % str(len(media_files_r)), indent=2)
            echo('Converted RAW files: %s' % str(len(media_files_c)), indent=2)

        if self.notify:
            macos_notify(title='Rename/Convert', message='Rename/Convert completed successfully!')

        progrun_update(
            name='Photo.rename_convert',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res=res)

# FIXME: res={}
class PhotoDatabaseRefresh(object):
    """-------------------------------------------------------------------------
    Script Name: Photo Database Refresh
    Script Author: Andoni Sooklaris
    Date: 2019-02-27

    Maintain a PostgreSQL database with EXIF metadata from photo library.
    Insert new photo/video records, update media files that have changed,
    and delete media files out of database if removed on filesystem.

    Keyword Arguments:
        rscript {str} -- Path to "photo_database_refresh.R" script (default: {None})
        dir {str} -- Path to photos dir to refresh database with (default: {None})
        pg_user {str} -- Postgres username (default: {None})
        pg_db {str} -- Postgres database name (default: {None})
        pg_schema {str} -- Postgres schema name (default: {None})
        logfile {str} -- SQL file to write results to. If None (default) do not write SQL log (default: {None})
        skip_insert {bool} -- Do not insert any records to Postgres database (default: {None})
        skip_update {bool} -- Do not update any records in Postgres database (default: {None})
        skip_delete {bool} -- Do not delete any records from Postgres database (default: {None})
        verbose {bool} -- Print messages to console (default: {None})
    -------------------------------------------------------------------------"""

    def __init__(
        self,
        rscript=None,
        dir=None,
        pg_user=None,
        pg_db=None,
        pg_schema=None,
        logfile=None,
        skip_insert=None,
        skip_update=None,
        skip_delete=None,
        verbose=None,
        ):
        self.rscript = rscript
        self.dir = dir
        self.pg_user = pg_user
        self.pg_db = pg_db
        self.pg_schema = pg_schema
        self.logfile = logfile
        self.skip_insert = skip_insert
        self.skip_update = skip_update
        self.skip_delete = skip_delete
        self.verbose = verbose

        assert isfile(self.rscript)
        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):
        # If parameter(s) unspecified, default to default in R script
        prog_start = datetime.now()
        cmdargs = []
        for param, value in self.params.items():
            if param == 'rscript':
                continue
            if value is not None:
                if type(value) in [int, float]:
                    cmdargs.append('--{}={}'.format(param, value))
                elif type(value) == bool:
                    if value:
                        cmdargs.append('--{}'.format(param))
                else:
                    cmdargs.append('--{}="{}"'.format(param, value))

        cmdargs = ' '.join(cmdargs)
        cmd = 'Rscript "{}" {}'.format(self.rscript, cmdargs)
        system(cmd)

        progrun_update(
            name='Photo.photo_database_refresh',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res={})  # FIXME: Add 'res' from Rscript, likely save as CSV then read with Python

class OrganizeWebsiteExportedPlacesPhotos(object):
    """
    Rename and organize photos exported from Lightroom for andonisooklaris.com.

    Arguments:
        exported_dir {str} -- Path to exported photo directory
        places_dir {str} -- Path to andonisooklaris.com Media/Places folder
    """

    def __init__(self, exported_dir, places_dir):
        self.exported_dir = exported_dir
        self.places_dir = places_dir

        assert isdir(self.exported_dir)
        assert isdir(self.places_dir)

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def split_photo_fname(self, fname):
        """
        Extract collection, date, title and location from filename.

        Arguments:
            fname {str} -- filename to split

        Returns:
            {tuple}
        """

        fbasename, ext = splitext(fname)
        collection, date, title, location = fbasename.split('_')

        # Get appended number by Lightroom upon export in case of identical image names
        if re.search(r'-\d+$', location):
            suffix = re.search(r'-\d+$', location).group(0)
            location = location.replace(suffix, '')
        else:
            suffix = ''

        return (collection, date, title, location, suffix, ext)

    def run(self):

        prog_start = datetime.now()

        chdir(self.exported_dir)
        fnames = listfiles()

        if not len(fnames):
            echo('No files to process!', abort=True)

        for fname in fnames:
            chdir(self.places_dir)
            collection, date, title, location, suffix, ext = self.split_photo_fname(fname)

            # Determine directory to move photo to in self.places_dir,
            # based on `collection`
            collection = collection.replace('Places-', '').split('-')
            collection = collection[0]
            dest_dir = join(self.places_dir, collection)
            if not isdir(dest_dir):
                mkdir(dest_dir)

            if re.search(r'^\d{8}$', date):
                date = date[0:4] + '-' + date[4:6] + '-' + date[6:]

            if title > '':
                newfname = date + ' ' + title + ' ' + location + suffix + ext
            else:
                newfname = date + ' ' + location + suffix + ext

            newfname = join(dest_dir, newfname)

            chdir(self.exported_dir)
            if not isfile(newfname):
                rename(fname, newfname)

        progrun_update(
            name='Photo.organize_website_exported_places_photos',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res={'fnames': fnames})


#!/usr/bin/python

import cv2
import numpy as np
import sys

# img1 = cv2.imread('maskout.bmp')
# img1 = img1.astype(np.float32)
# img2 = cv2.imread('maskout_-15.bmp')
# img2 = img2.astype(np.float32)

img1 = cv2.imread('IMG_7094_half.JPG')
img1 = img1.astype(np.float32)
img2 = cv2.imread('screen_1920x1080_4.png')
img2 = img2.astype(np.float32)

mapper = cv2.reg_MapperGradAffine()
mappPyr = cv2.reg_MapperPyramid(mapper)

resMap = mappPyr.calculate(img1, img2)
mapShift = cv2.reg.MapTypeCaster_toAffine(resMap)

print(mapShift.getLinTr())
print(mapShift.getShift())


"""
Modified implementation of Snurr Group's EQeq method for generation partial
atomic charges.
"""

from __future__ import print_function, division

import argparse
import sys

import numpy
from numpy import linalg

from main import chemistry, config, inputdata
from reference import ionization

# Define a fundamental constant
# K = (electronic charge ^2) / ( 4 * pi * vacuum permittivity)
K = 14.4

class EQeq:
    """Calculate partial charges from preprocessed input data"""

    def __init__(self, h0_value, s_charge):
        self.h0_value = -h0_value
        self.s_charge = s_charge

    def ip_list(self, element):
        """
        Read in the list of ionization potentials for an element from
        reference data.
        :type element: chemistry.Element
        """
        # Return ip_list immediately if cached.
        # REMEMBER that if self.h0_value is being optimised, this hasattr()
        # must be commented out since elements are persistent across steps.
        if hasattr(element, "ip_list"):
            return element.ip_list

        if element.atomic_number >= len(ionization.data):
            print("Error: No IP data available for", element.symbol)
            sys.exit()

        if element.atomic_number == 0:
            ip_list = [self.h0_value, 13.598]
        else:
            ip_list = ionization.data[element.atomic_number][1:]

        element.ip_list = ip_list
        return ip_list

    def electronegativity(self, element):
        """
        Calculate the Mulliken electronegativity from ionization potential
        data.
        :rtype : float
        :type element: chemistry.Element
        """
        if hasattr(element, "electronegativity"):
            return element.electronegativity

        ip_list = self.ip_list(element)

        x_0 = ip_list[0]
        x_1 = ip_list[1]
        x = 0.5 * (x_1 + x_0)

        element.electronegativity = x
        return x

    def chemical_hardness(self, element):
        """
        Calculate the chemical hardness / idempotential from ionization
        potential data.
        :rtype : float
        :type element: chemistry.Element
        """
        if hasattr(element, "chemical_hardness"):
            return element.chemical_hardness

        ip_list = self.ip_list(element)
        j_0 = ip_list[0]
        j_1 = ip_list[1]
        j = j_1 - j_0

        element.chemical_hardness = j
        return j

    def j_ij(self, atom1, atom2):
        """
        Calculate J values including scaling for distance

        :rtype : float
        :type atom1: chemistry.Atom
        :type atom2: chemistry.Atom
        """
        if atom1 == atom2:
            return self.chemical_hardness(atom1.element)
        else:
            j_i = self.chemical_hardness(atom1.element)
            j_j = self.chemical_hardness(atom2.element)

            r_ab = atom1.separation(atom2)

            # Nishimoto-Mataga
            j_ab = 1 / (r_ab + (2 / (j_i + j_j)))

            j_ab *= (K / 2)

            return j_ab

    def partial_charges(self, atoms, total_charge):
        """Calculate the partial atomic charges for every atom
        :type total_charge: int
        :type atoms: list
        """

        def construct_j_matrix(atom_list):
            """
            Take the list of atom objects and return the matrix of J
            (and modified J) values.
            :rtype : scipy.array
            :type atom_list: list
            """
            matrix_lh = []
            matrix_rh = []

            for row in range(len(atom_list) - 1):
                temp_row = []
                for column in range(len(atom_list)):
                    temp_row.append(self.j_ij(atom_list[row],
                                              atom_list[column]))
                matrix_lh.append(temp_row)

            for row in range(1, len(atom_list)):
                temp_row = []
                for column in range(len(atom_list)):
                    temp_row.append(self.j_ij(atom_list[row],
                                              atom_list[column]))
                matrix_rh.append(temp_row)

            matrix_lh = numpy.array(matrix_lh)
            matrix_rh = numpy.array(matrix_rh)

            ones_row = numpy.array([[1 for x in range(0, len(atom_list))]])
            j_matrix = numpy.subtract(matrix_lh, matrix_rh)
            final_matrix = numpy.concatenate((ones_row, j_matrix))

            return final_matrix

        def construct_x_matrix(atom_list, charge):
            """
            Take the list of atom objects and return the matrix of X values.
            :type atom_list: list
            :type charge: int
            :rtype : scipy.array
            """
            x_matrix = numpy.array([[charge]])
            for x in range(1, len(atom_list)):
                element0 = atom_list[x - 1].element
                element1 = atom_list[x].element
                difference = self.electronegativity(element1) - \
                             self.electronegativity(element0)
                x_matrix = numpy.concatenate((x_matrix, [[difference]]))
            return x_matrix

        # atom_lists = [[] for _ in range(atoms[-1].molecule + 1)]
        atom_lists = [[]]

        for atom in atoms:
            # atom_lists[atom.molecule].append(atom)
            atom_lists[0].append(atom)

        partials = []
        for molecule in atom_lists:
            j_matrix = construct_j_matrix(molecule)
            x_matrix = construct_x_matrix(molecule, total_charge)
            q_matrix = linalg.solve(j_matrix, x_matrix)

            for x in range(len(molecule)):
                partial_charge = q_matrix[x][0] * self.s_charge
                partials.append(partial_charge)
                molecule[x].partialcharge = partial_charge

        return partials

    @staticmethod
    def system_energy(pairs):
        """
        Return the sum of the pairwise electrostatic interactions between
        all atoms in the system.
        :type pairs: list
        :rtype : float
        """

        def pairwise_term(pair_in):
            """
            Return the coulombic energy between two point charges.

            :type pair_in: chemistry.AtomPair
            :rtype : float
            """
            q1 = pair_in.atom1.partialcharge
            q2 = pair_in.atom2.partialcharge

            r_ab = pair_in.separation()

            term_out = (q1 * q2) / r_ab

            return term_out

        total = 0
        for pair in pairs:
            if pair.atom1.molecule != pair.atom2.molecule:
                term = K * pairwise_term(pair)
                total += term
        energy = total * chemistry.ev_to_kcal

        return energy

def run_program(file_in, h0_value=1.839111328, s_charge=4.834350586):
    """
    Run the program and return the total electrostatic energy

    :type file_in: str
    :type h0_value: float
    :type s_charge: float
    :rtype : float
    """

    data = inputdata.InputData(file_in)
    data.get_all()

    eqeq = EQeq(h0_value, s_charge)
    eqeq.partial_charges(data.atoms(), data.total_charge_value)

    energy = eqeq.system_energy(data.pairs())

    return energy

if __name__ == "__main__":

    eqeq_description = "Calculate EQeq Electrostatic interaction energy"
    verbose_help = "Print all data verbosely."
    parameters_help = "Manually specify the parameters used in EQeq."

    parser = argparse.ArgumentParser(description=eqeq_description)
    parser.add_argument("-v", "--verbose", action="store_true",
                        help=verbose_help)
    parser.add_argument("-p", "--parameters", type=float, nargs=2,
                        default=False, help=parameters_help,
                        metavar=('h0_value', 's_charge'))
    parser.add_argument("file", type=str, nargs="+",
                        help="A Gaussian input file.")
    args = parser.parse_args()

    if args.verbose:
        config.verbose = True

    for filename in args.file:
        if args.parameters:
            print(run_program(filename, *args.parameters))
        else:
            print(run_program(filename))


import logging
import re
from steelfig.models.mysql import Mysql
from steelfig.models.account import Account

logger = logging.getLogger(__name__)
class Wishlist(object):

    def __init__(self, event_id=None, account_id=None):
        self.id = None
        self.event_id = event_id
        self.account_id = account_id

        if event_id and account_id:
            self.id = self.get_id_from_attendee(event_id, account_id)

    def get_id_from_attendee(self, event_id, account_id):
        db = Mysql()
        cursor = db.get_cursor()
        cursor.execute('''
            SELECT wishlists.id
            FROM wishlists
            JOIN attendees ON wishlists.attendee_id = attendees.id
            WHERE attendees.event_id = %s
            AND attendees.account_id = %s
            LIMIT 1
        ''', (event_id, account_id))
        data = cursor.fetchone()
        return False if not data else int(data.get('id'))

    def get_my_items(self):
        data = []
        if self.id:
            db = Mysql()
            cursor = db.get_cursor()
            cursor.execute('''
                SELECT id, name, url, price, sort, comments, created_at, updated_at
                FROM wishlist_items
                WHERE wishlist_id = %s
                ORDER BY sort ASC
            ''', (self.id,))
            data = cursor.fetchall()
        return {'event_wishlist_items': data}

    def get_buddies_items(self):
        data = []
        if self.id:
            db = Mysql()
            cursor = db.get_cursor()
            cursor.execute('''
                SELECT id, name, url, price, sort, comments, is_purchased, created_at, updated_at
                FROM wishlist_items
                WHERE wishlist_id = %s
                ORDER BY sort ASC
            ''', (self.id,))
            data = cursor.fetchall()
        return {'event_wishlist_items': data}

    def set_item(self, account_id, params, id=None):
        db = Mysql()
        cursor = db.get_cursor()
        if not params.get('name'):
            raise Exception('no name found')

        # Convert money to integer to avoid floating point crazyness
        price = str(params.get('price'))
        if price:
            price = re.sub(r'[^0-9\.]', '', price)
            price = "{:10.2f}".format(float(price))
            price = float(price) * 100

        # If url doesn't start with http, prepend it
        url = params.get('url')
        if url and not url.startswith('http'):
            url = 'http://' + url

        cursor.execute('''
            CALL set_wishlist_item(
                if(%s = '', NULL, %s),
                if(%s = '', NULL, %s),
                if(%s = '', NULL, %s),
                %s, %s, %s,
                if(%s = '', NULL, %s),
                %s, %s, %s)
        ''', (id or '', id or '',
              account_id, account_id,
              params.get('event_id', ''), params.get('event_id'),
              params.get('name', ''),
              params.get('barcode', ''),
              url,
              price, price,
              params.get('is_purchased', 0),
              params.get('sort', 0),
              params.get('comments')))
        data = cursor.fetchone()
        return {'event_wishlist_item': data}

    def delete_item(self, account_id, wishlist_item_id):
        db = Mysql()
        cursor = db.get_cursor()
        cursor.execute('''
            DELETE FROM wishlist_items
            WHERE id = %s
            AND wishlist_id IN (
                SELECT wishlists.id
                FROM wishlists
                JOIN attendees ON wishlists.attendee_id = attendees.id
                WHERE attendees.account_id = %s)
            LIMIT 1
        ''', (wishlist_item_id, account_id))
        return True


from discord.ext import commands
from collections import Counter
from utils import checks
from utils.db.datatype import *

import logging
import discord
import datetime
import psutil

from config import LOGGING_CHANNEL

client = discord.Client()
log = logging.getLogger(__name__)

class ActionReason(commands.Converter):
    async def convert(self, ctx, argument):
        ret = f'{ctx.author} (ID: {ctx.author.id}): {argument}'

        if len(ret) > 512:
            reason_max = 512 - len(ret) - len(argument)
            raise commands.BadArgument(f'reason is too long ({len(argument)}/{reason_max})')
        return ret

class BannedMember(commands.Converter):
    async def convert(self, ctx, argument):
        ban_list = await ctx.guild.bans()

        try:
            member_id = int(argument, base=10)
            entity = discord.utils.find(lambda u: u.user.id == member_id, ban_list)
        except ValueError:
            entity = discord.utils.find(lambda u: str(u.user.name) == argument, ban_list)

        if entity is None:
            raise commands.BadArgument("Not a valid previously-banned member.")
        return entity

class utility:

    def __init__(self, bot):
        self.bot = bot
        self.process = psutil.Process()

    @commands.command(hidden=False)
    async def ping(self, ctx):
        await ctx.send('pong !')

    @commands.command(hidden=False)
    async def back(self, ctx):
        '''Set an afk reason'''
        afk_back(str(ctx.guild.id), ctx.author.name)
        await ctx.send(f"{ctx.author.mention} is back !")

    @commands.command(hidden=False)
    async def afk(self, ctx):
        reason = "Afk"
        if len(ctx.message.content.split(' ')[1:]):
            reason = " ".join(ctx.message.content.split(' ')[1:])
        afk_reason(str(ctx.guild.id), ctx.author.name, reason)
        await ctx.send(f"{ctx.author.name} is afk :\"{reason}\"")

    @client.event
    async def on_message(self, message: discord.Message):
        if not message.guild:
            return
        if message.author == message.guild.me:
            return
        for mention in message.mentions:
            reason = is_afk(message.guild.id, mention.name)
            if reason:
               await message.channel.send( f'{mention.name} is AFK : {reason}')

    @commands.command()
    @commands.guild_only()
    @checks.has_permissions(kick_members=True)
    async def kick(self, ctx, member: discord.Member, *, reason: ActionReason = None):
        """Kicks a member from the server.

        In order for this to work, the bot must have Kick Member permissions.

        To use this command you must have Kick Members permission.
        """

        if reason is None:
            reason = f'Action done by {ctx.author} (ID: {ctx.author.id})'

        await member.kick(reason=reason)
        await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    @commands.guild_only()
    @checks.has_permissions(ban_members=True)
    async def ban(self, ctx, member: discord.Member, *, reason: ActionReason = None):
        """Ban a member from the server.

        In order for this to work, the bot must have Ban Member permissions.

        To use this command you must have Ban Members permission.
        """

        if reason is None:
            reason = f'Action done by {ctx.author} (ID: {ctx.author.id})'

        await member.ban(reason=reason)
        await ctx.send('\N{OK HAND SIGN}')

    @commands.command()
    @commands.guild_only()
    @checks.has_permissions(ban_members=True)
    async def unban(self, ctx, member: BannedMember, *, reason: ActionReason = None):
        """Unbans a member from the server.

        You can pass either the ID of the banned member or the Name#Discrim
        combination of the member. Typically the ID is easiest to use.

        In order for this to work, the bot must have Ban Member permissions.

        To use this command you must have Ban Members permissions.
        """

        if reason is None:
            reason = f'Action done by {ctx.author} (ID: {ctx.author.id})'

        await ctx.guild.unban(member.user, reason=reason)
        if member.reason:
            await ctx.send(f'Unbanned {member.user} (ID: {member.user.id}), previously banned for {member.reason}.')
        else:
            await ctx.send(f'Unbanned {member.user} (ID: {member.user.id}).')

    async def _full_cleanup_strategy(self, ctx, search):
        prefixes = tuple('!') # thanks startswith

        def check(m):
            return m.pinned == False

        deleted = await ctx.channel.purge(limit=search, check=check, before=ctx.message)
        return Counter(m.author.display_name for m in deleted)

    @commands.command()
    @checks.has_permissions(manage_messages=True)
    async def clear(self, ctx, search=100):
        """Cleans up the bot's messages from the channel.

        If a search number is specified, it searches that many messages to delete.
        If the bot has Manage Messages permissions then it will try to delete
        messages that look like they invoked the bot as well.

        After the cleanup is completed, the bot will send you a message with
        which people got their messages deleted and their count. This is useful
        to see which users are spammers.

        You must have Manage Messages permission to use this.
        """

        strategy = self._full_cleanup_strategy

        spammers = await strategy(ctx, search)
        deleted = sum(spammers.values())
        messages = [f'{deleted} message{" was" if deleted == 1 else "s were"} removed.']
        if deleted:
            messages.append('')
            spammers = sorted(spammers.items(), key=lambda t: t[1], reverse=True)
            messages.extend(f'- **{author}**: {count}' for author, count in spammers)

        await ctx.message.delete()
        await ctx.send('\n'.join(messages), delete_after=10)

async def on_error(self, event, *args, **kwargs):
    e = discord.Embed(title='Event Error', colour=0xa32952)
    e.add_field(name='Event', value=event)
    e.description = f'```py\n{traceback.format_exc()}\n```'
    e.timestamp = datetime.datetime.utcnow()
    ch = self.get_channel(LOGGING_CHANNEL)
    try:
        await ch.send(embed=e)
    except:
        pass

def setup(bot):
    if not hasattr(bot, 'command_stats'):
        bot.command_stats = Counter()

    if not hasattr(bot, 'socket_stats'):
        bot.socket_stats = Counter()

    bot.add_cog(utility(bot))
    commands.Bot.on_error = on_error


import torch
import torch.nn.functional as F
from libs import InPlaceABN, InPlaceABNSync
from torch import nn
from torch.nn import init
import math

class _NonLocalNdGc(nn.Module):

    def __init__(self, dim, inplanes, planes, downsample, use_gn, lr_mult, use_out, add_conv):
        assert dim in [1, 2, 3], "dim {} is not supported yet".format(dim)
        assert add_conv in [None, 'one_fc', 'two_fc']
        if dim == 3:
            conv_nd = nn.Conv3d
            if downsample:
                max_pool = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
            else:
                max_pool = None
            bn_nd = nn.BatchNorm3d
        elif dim == 2:
            conv_nd = nn.Conv2d
            if downsample:
                max_pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
            else:
                max_pool = None
            bn_nd = nn.BatchNorm2d
        else:
            conv_nd = nn.Conv1d
            if downsample:
                max_pool = nn.MaxPool1d(kernel_size=2, stride=2)
            else:
                max_pool = None
            bn_nd = nn.BatchNorm1d

        super(_NonLocalNdGc, self).__init__()
        self.conv_query = conv_nd(inplanes, planes, kernel_size=1)
        self.conv_key = conv_nd(inplanes, planes, kernel_size=1)
        self.conv_mask = conv_nd(inplanes, 1, kernel_size=1)
        if add_conv == 'one_fc':
            self.channel_add_conv=nn.Sequential(
                nn.Conv2d(inplanes, inplanes, kernel_size=1),
                nn.LayerNorm([inplanes, 1, 1]))
        elif add_conv == 'two_fc':
            self.channel_add_conv = nn.Sequential(
                nn.Conv2d(inplanes, inplanes//4, kernel_size=1),
                nn.LayerNorm([inplanes//4, 1, 1]),
                nn.ReLU(inplace=True),  # yapf: disable
                nn.Conv2d(inplanes//4, inplanes, kernel_size=1))
        if use_out:
            self.conv_value = conv_nd(inplanes, planes, kernel_size=1)
            self.conv_out = conv_nd(planes, inplanes, kernel_size=1, bias=False)
        else:
            self.conv_value = conv_nd(inplanes, inplanes, kernel_size=1, bias=False)
            self.conv_out = None
        self.softmax = nn.Softmax(dim=2)
        self.downsample = max_pool
        # self.norm = nn.GroupNorm(num_groups=32, num_channels=inplanes) if use_gn else InPlaceABNSync(num_features=inplanes)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.scale = math.sqrt(planes)
        self.add_conv = add_conv

        self.reset_parameters()
        self.reset_lr_mult(lr_mult)

    def reset_parameters(self):

        for m in self.modules():
            if isinstance(m, nn.Conv3d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    init.zeros_(m.bias)
                m.inited = True
        #init.constant_(self.norm.weight, 0)
        #init.constant_(self.norm.bias, 0)
        #self.norm.inited = True

    def reset_lr_mult(self, lr_mult):
        if lr_mult is not None:
            for m in self.modules():
                m.lr_mult = lr_mult
        else:
            print('not change lr_mult')

    def forward(self, x):
        # [N, C, T, H, W]
        residual = x
        # [N, C, T, H', W']
        if self.downsample is not None:
            input_x = self.downsample(x)
        else:
            input_x = x

        # [N, C', T, H, W]
        query = self.conv_query(x)
        # [N, C', T, H', W']
        key = self.conv_key(input_x)
        value = self.conv_value(input_x)
        # [N, 1, H', W']
        mask = self.conv_mask(input_x)

        # [N, C', H x W]
        query = query.view(query.size(0), query.size(1), -1)
        # [N, C', H' x W']
        key = key.view(key.size(0), key.size(1), -1)
        value = value.view(value.size(0), value.size(1), -1)
        # [N, 1, H'x W']
        mask = mask.view(mask.size(0), mask.size(1), -1)

        # [N, T x H x W, T x H' x W']
        sim_map = torch.bmm(query.transpose(1, 2), key)
        sim_map = sim_map/self.scale
        sim_map = self.softmax(sim_map)
        # [N, 1, H'x W']
        mask = self.softmax(mask)

        # [N, T x H x W, C']
        out_sim = torch.bmm(sim_map, value.transpose(1, 2))
        # [N, C', T x H x W]
        out_sim = out_sim.transpose(1, 2)
        # [N, C', T,  H, W]
        out_sim = out_sim.view(out_sim.size(0), out_sim.size(1), *x.size()[2:])
        # [N, C, T,  H, W]
        if self.conv_out is not None:
            out_sim = self.conv_out(out_sim)
        # if self.norm is not None:
        #     out = self.norm(out)
        out_sim = self.gamma * out_sim

        # [N, C', 1, 1]
        out_gc = torch.bmm(value, mask.permute(0,2,1)).unsqueeze(-1)

        if self.add_conv is not None:
            out_gc = self.channel_add_conv(out_gc)

        out = residual + out_sim + out_gc
        return out

class NonLocal2dGc(_NonLocalNdGc):

    def __init__(self, inplanes, planes, downsample=False, use_gn=False, lr_mult=None, use_out=False, add_conv=None):
        super(NonLocal2dGc, self).__init__(dim=2, inplanes=inplanes, planes=planes, downsample=downsample, use_gn=use_gn, lr_mult=lr_mult, use_out=use_out, add_conv=add_conv)

class NonLocal3dGc(_NonLocalNdGc):

    def __init__(self, inplanes, planes, downsample, use_gn, lr_mult, use_out):
        super(NonLocal3dGc, self).__init__(dim=3, inplanes=inplanes, planes=planes, downsample=downsample, use_gn=use_gn, lr_mult=lr_mult, use_out=use_out)


from time import time

from legal_doc_processing import logger

from legal_doc_processing.utils import get_pipeline, get_spacy
from legal_doc_processing.base.base import Base
from legal_doc_processing.press_release.utils import press_release_X_y

class PressRelease(Base):
    """main press release doc class """

    def __init__(
        self,
        text: str,
        source: str,
        nlpipe=None,
        nlspa=None,
        n_lines: int = 6,
    ):

        Base.__init__(
            self,
            text=text,
            obj_name="PressRelease",
            source=source,
            nlpipe=nlpipe,
            nlspa=nlspa,
        )

        # set all
        self.set_all()

def press_release_df(
    juridiction="", nlspa=None, nlpipe=None, sample=0.25, max_init_time=3.0
):
    """ """

    assert juridiction in ["cftc", "cfbp", "doj", "sec", ""]

    max_init_time = 3.0

    # load
    if not nlpipe:
        nlpipe = get_pipeline()
    if not nlspa:
        nlspa = get_spacy()

    # dataframe
    df = press_release_X_y(juridiction=juridiction, sample=sample)

    # Press Releae

    juri = juridiction.lower().strip()
    if juridiction:
        # selec juridiction
        select_jur = lambda i: str(i).lower().strip() == juri
        df = df.loc[df.juridiction.apply(select_jur), :]
        # make pr
        make_pr = lambda i: PressRelease(i, source=juri, nlpipe=nlpipe, nlspa=nlspa)
        df["pr"] = df.press_release_text.apply(make_pr)
    else:
        # make pr
        make_pr = lambda i, j: PressRelease(i, source=j, nlpipe=nlpipe, nlspa=nlspa)
        df["pr"] = [make_pr(i, j) for i, j in zip(df.press_release_text, df.juridiction)]

    return df

# def from_file(file_path, source, nlpipe=None, nlspa=None):
#     return base_from_file(file_path, source, PressRelease, nlpipe=nlpipe, nlspa=nlspa)

# def from_text(txt, source, nlpipe=None, nlspa=None):
#     return base_from_text(txt, source, PressRelease, nlpipe=nlpipe, nlspa=nlspa)

# def from_url(txt, source, nlpipe=None, nlspa=None):
#     return base_from_url(txt, source, PressRelease, nlpipe=nlpipe, nlspa=nlspa)

class _PressRelease:
    PressRelease = PressRelease
    load_X_y = press_release_X_y
    load_df = press_release_df