from math import ceil
from typing import Callable, Dict, Iterable, List, Tuple

import torch
import torch.nn.functional as F
from ppq.core import QuantizationProperty, QuantizationStates
from ppq.executor import BaseGraphExecutor
from ppq.executor.base import OPERATION_FORWARD_TABLE
from ppq.IR import BaseGraph, Operation, QuantableOperation, Variable
from ppq.IR.search import Path, SearchableGraph, TraversalCommand
from ppq.log import NaiveLogger
from ppq.quantization.measure import torch_mean_square_error
from ppq.quantization.observer import CalibrationHook, OperationObserver
from ppq.quantization.observer.range import TorchHistObserver
from ppq.quantization.qfunction import (BaseQuantFunction,
                                        PPQLinearQuantFunction)
from tqdm import tqdm

from .base import QuantizationOptimizationPass

logger = NaiveLogger.get_logger('PPQ')

OPTIMIZATION_LAYERTYPE_CONFIG = {
    1: {'Relu', 'MaxPool', 'GlobalMaxPool', 'PRelu', 'AveragePool', 'GlobalAveragePool'}, # support more relay operation types
}
EQUALIZATION_OPERATION_TYPE = {'Conv', 'Gemm', 'ConvTranspose'} # support all computing op types

class SSDEqualizationPass(QuantizationOptimizationPass):
    """PPQ Custimized Layerwise Equalization Pass.

    This is another layerwise equalization pass which takes quantization error into consideration, for more
    details of equalization, please refer to LayerwiseEqualizationPass

    Equalization algo will only take effect when the estimated loss is less than loss_threshold * original_loss

    Compared with LayerwiseEqualizationPass, it guarantees a better result than the original model while consuming
    more time for loss estimation
    """
    def __init__(
        self,
        optimize_level: int=1,
        channel_ratio: float = 0.5,
        loss_threshold: float = 0.8,
        layer_norm: bool = False,
        quant_func: BaseQuantFunction = PPQLinearQuantFunction,
        iteration: int = 3
    ):
        """SSD Equalization Pass With Loss Checking.

        Args:
            optimize_level (int, optional): level of optimization. Only support level 1 for now.

            channel_ratio (float, optional): all values below this ratio of maximum value of corresponding weight will
            be clipped to this ratio when calculating scale. Defaults to 0.5.

            loss_threshold (float, optional): the estimated loss should be below this threshold * original_loss for algo
            to take effect. Defaults to 0.8.

            layer_norm (bool, optional): whether to apply weight normalization. Defaults to False.

            quant_func (BaseQuantFunction, optional): quantization function. Defaults to PPQLinearQuantFunction.

            iteration (int, optional): num of iterations to run, usually 3 would be enough. Defaults to 3.
        """
        self.channel_ratio   = channel_ratio
        self.loss_threshold  = loss_threshold
        self.layer_norm      = layer_norm
        self.quant_func      = quant_func
        self.start_op_types  = EQUALIZATION_OPERATION_TYPE
        self.relay_op_types  = OPTIMIZATION_LAYERTYPE_CONFIG[optimize_level]
        self.end_op_types    = EQUALIZATION_OPERATION_TYPE
        self.iteration       = iteration
        super().__init__(name='SSD Equalization Pass')

    def collect_all_pairs(self, graph: BaseGraph) -> List[List[Operation]]:
        def limitation(path: Path) -> bool:
            for op in path.tolist()[:-1]:
                if len(graph.get_downstream_operations(op)) != 1:
                    return False
            return True

        search_engine = SearchableGraph(graph)
        forward_matchings = search_engine.process(TraversalCommand(
            sp_expr=lambda x: x.type in self.start_op_types,
            rp_expr=lambda x, y: y.type in self.relay_op_types,
            ep_expr=lambda x: x.type in self.end_op_types,
            direction='down'))
        return [path.tolist() for path in forward_matchings if limitation(path)]

    def collect_activation_range(
        self,
        pair:  List[Operation],
        executor: BaseGraphExecutor,
        data_loader: Iterable,
        collate_fn: Callable,
        calib_steps: int
    ) -> Dict[Operation, torch.Tensor]:
        """Collect activation ranges for Conv ops in the pair.

        Args:
            pair (List[Operation]): equalization pair
            executor (BaseGraphExecutor): graph executor
            data_loader (Iterable): data loader
            collate_fn (Callable): batch func
            calib_steps (int): calibration steps

        Returns:
            Dict[Operation, torch.Tensor]: activation range of Conv ops in the pair
        """
        op_act_ranges = {op: 0.0 for op in pair if op.type in self.start_op_types}
        output_names = [op.outputs[0].name for op in op_act_ranges]
        calib_step = 0
        for calib_epoch in range(ceil(calib_steps / len(data_loader))):
            for _, data in enumerate(data_loader):
                if collate_fn is not None:
                    data = collate_fn(data)
                outputs = executor.forward(inputs=data, output_names=output_names)
                for op in op_act_ranges:
                    op_conv_data = outputs[output_names.index(op.outputs[0].name)]
                    op_conv_data_relu = F.relu(op_conv_data) # take abs value
                    op_conv_data_relu = op_conv_data_relu.permute(
                        1, 0, *(range(op_conv_data_relu.ndim)[2:])).contiguous()
                    op_conv_data_relu = op_conv_data_relu.reshape((op_conv_data_relu.shape[0], -1))
                    op_conv_data_relu = op_conv_data_relu.max(1)[0]
                    op_act_ranges[op] += op_conv_data_relu
                calib_step += 1
                if calib_step >= calib_steps: break
        for op in op_act_ranges:
            op_act_ranges[op] /= calib_steps

        for op,op_range in op_act_ranges.items():
            assert isinstance(op_range, torch.Tensor)
            op_range_max = op_range.max()
            adjust_range = torch.where(op_range < op_range_max * self.channel_ratio, op_range_max * self.channel_ratio, op_range)
            op_act_ranges[op] = adjust_range
        return op_act_ranges

    def layer_weight_norm(self, pairs: List[List[Operation]]):
        for pair in pairs:
            op_first_weight = pair[0].parameters[0].value
            op_second_weight = pair[-1].parameters[0].value
            op_first_max = op_first_weight.abs().max()
            op_second_max = op_second_weight.abs().max()
            scale = (op_second_max / op_first_max).sqrt()
            pair[0].parameters[0].value = pair[0].parameters[0].value * scale
            if len(pair[0].parameters) > 1:
                pair[0].parameters[1].value = pair[0].parameters[1].value * scale
            pair[-1].parameters[0].value = pair[-1].parameters[0].value / scale

    def prepare_weight_for_equalization(self, pair: List[Operation]) -> Tuple[torch.Tensor]:
        first_computing_op_weight = pair[0].parameters[0].value
        last_computing_op_weight  = pair[-1].parameters[0].value

        assert isinstance(first_computing_op_weight, torch.Tensor)
        assert isinstance(last_computing_op_weight, torch.Tensor)

        if pair[0].type == 'Conv':
            # [C_out, C_in, K, K]
            C_out = first_computing_op_weight.shape[0]
            first_weight_range = first_computing_op_weight.reshape(C_out, -1).abs().max(dim=1)[0]

        elif pair[0].type == 'Gemm':
            if pair[0].attributes.get('transB', 0):
                # [C_out, C_in]
                first_weight_range = first_computing_op_weight.abs().max(dim=1)[0]
            else:
                # [C_in, C_out]
                first_weight_range = first_computing_op_weight.abs().max(dim=0)[0]

        elif pair[0].type == 'ConvTranspose':
            # [C_in, C_out // g, K, K]
            num_group = pair[0].attributes.get('group', 1)
            C_in, C_out_g, K1, K2 = first_computing_op_weight.shape
            first_computing_op_weight = first_computing_op_weight.reshape(num_group, C_in // num_group, C_out_g, K1, K2)
            first_computing_op_weight = first_computing_op_weight.permute(0, 2, 1, 3, 4).contiguous()
            first_computing_op_weight = first_computing_op_weight.reshape(num_group * C_out_g, -1)
            first_weight_range = first_computing_op_weight.abs().max(dim=1)[0]

        if pair[-1].type == 'Conv':
            # [C_out, C_in // g, K, K]
            num_group = pair[-1].attributes.get('group', 1)
            C_out, C_in_g, K1, K2 = last_computing_op_weight.shape
            last_computing_op_weight = last_computing_op_weight.reshape(num_group, C_out // num_group, C_in_g, K1, K2)
            last_computing_op_weight = last_computing_op_weight.permute(0, 2, 1, 3, 4).contiguous()
            last_computing_op_weight = last_computing_op_weight.reshape(num_group * C_in_g, -1)
            last_weight_range = last_computing_op_weight.abs().max(dim=1)[0]

        elif pair[-1].type == 'Gemm':
            C_out = first_weight_range.shape[0]
            if pair[-1].attributes.get('transB', 0):
                if C_out != last_computing_op_weight.shape[1]:
                    last_computing_op_weight = last_computing_op_weight.reshape(last_computing_op_weight.shape[0], C_out, -1)
                    last_computing_op_weight = last_computing_op_weight.permute(1, 0, 2).contiguous().reshape(C_out, -1)
                    last_weight_range = last_computing_op_weight.abs().max(dim=1)[0]
                else:
                    last_weight_range = last_computing_op_weight.abs().max(dim=0)[0]
            else:
                if C_out != last_computing_op_weight.shape[0]:
                    last_computing_op_weight = last_computing_op_weight.reshape(C_out, -1)
                    last_weight_range = last_computing_op_weight.abs().max(dim=1)[0]
                else:
                    last_weight_range = last_computing_op_weight.abs().max(dim=1)[0]

        elif pair[-1].type == 'ConvTranspose':
            C_in = last_computing_op_weight.shape[0]
            last_weight_range = last_computing_op_weight.reshape(C_in, -1).abs().max(dim=1)[0]

        return first_weight_range, last_weight_range

    def write_back(self, pair: List[Operation], scale: torch.Tensor) -> None:
        first_computing_op_weight = pair[0].parameters[0].value
        last_computing_op_weight  = pair[-1].parameters[0].value

        assert isinstance(first_computing_op_weight, torch.Tensor)
        assert isinstance(last_computing_op_weight, torch.Tensor)

        if pair[0].type == 'Conv':
            pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(-1, 1, 1, 1)

        elif pair[0].type == 'Gemm':
            if pair[0].attributes.get('transB', 0):
                pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(-1, 1)
            else:
                pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(1, -1)

        elif pair[0].type == 'ConvTranspose':
            num_group = pair[0].attributes.get('group', 1)
            C_in, C_out_g, K1, K2 = first_computing_op_weight.shape
            first_computing_op_weight = first_computing_op_weight.reshape(num_group, C_in // num_group, C_out_g, K1, K2)
            first_computing_op_weight = first_computing_op_weight * scale.reshape(num_group, 1, -1, 1, 1)
            pair[0].parameters[0].value = first_computing_op_weight.reshape(C_in, C_out_g, K1, K2)

        if len(pair[0].parameters) > 1:
            pair[0].parameters[1].value = pair[0].parameters[1].value * scale

        if pair[-1].type == 'Conv':
            num_group = pair[-1].attributes.get('group', 1)
            C_out, C_in_g, K1, K2 = last_computing_op_weight.shape
            last_computing_op_weight = last_computing_op_weight.reshape(num_group, C_out // num_group, C_in_g, K1, K2)
            last_computing_op_weight = last_computing_op_weight / scale.reshape(num_group, 1, -1, 1, 1)
            pair[-1].parameters[0].value = last_computing_op_weight.reshape(C_out, C_in_g, K1, K2)

        elif pair[-1].type == 'Gemm':
            if pair[-1].attributes.get('transB', 0):
                if scale.numel() != last_computing_op_weight.shape[1]:
                    last_computing_op_weight = last_computing_op_weight.reshape(last_computing_op_weight.shape[0], scale.numel(), -1)
                    last_computing_op_weight = last_computing_op_weight / scale.reshape(1, -1, 1)
                    pair[-1].parameters[0].value = last_computing_op_weight.reshape(last_computing_op_weight.shape[0], -1)
                else:
                    pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(1, -1)
            else:
                if scale.numel() != last_computing_op_weight.shape[0]:
                    last_computing_op_weight = last_computing_op_weight.reshape(scale.numel(), -1, last_computing_op_weight.shape[-1])
                    last_computing_op_weight = last_computing_op_weight / scale.reshape(-1, 1, 1)
                    pair[-1].parameters[0].value = last_computing_op_weight.reshape(-1, last_computing_op_weight.shape[-1])
                else:
                    pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(-1, 1)

        elif pair[-1].type == 'ConvTranspose':
            pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(-1, 1, 1, 1)

    def one_step_equalization(
        self,
        pair: List[Operation],
        op_act_channel_range: Dict[Operation, torch.Tensor]={},
        algo_type: int=2,
        ssd_min_scale: float=8,
        ssd_max_scale: float=2,
        dfq_min_scale: float=0.1,
        dfq_max_scale: float=10,
        eps: float=1e-8
    ):
        """Equalization step with scale being calculated in the way specified
        by algo_type.

        Args:
            pair (List[Operation]): a list of operations representing a equalzation pair
            op_act_channel_range (Dict[Operation, torch.Tensor]): channel-wise activation range of all Conv ops in the graph
            algo_type (int, optional): minor algo type. 0 represents dfq algo, 1~3 represents ssd algo. Defaults to 2.
            ssd_min_scale (float, optional): minimum clip value of scale for ssd algo. Defaults to 8.
            ssd_max_scale (float, optional): maximum clip value of scale for ssd algo. Defaults to 2.
            dfq_min_scale (float, optional): minimum clip value of scale for dfq algo. Defaults to 0.1.
            dfq_max_scale (float, optional): maximum clip value of scale for dfq algo. Defaults to 10.
            eps (float, optional): small constant for numerical stability. Defaults to 1e-8.
        """
        first_weight_range, last_weight_range = self.prepare_weight_for_equalization(pair)

        if algo_type == 0:
            # dfq
            scale = torch.sqrt(last_weight_range / (first_weight_range + eps))
            scale = torch.clamp(scale, dfq_min_scale, dfq_max_scale)

        else:
            first_weight_range = torch.where(first_weight_range < first_weight_range.max() * self.channel_ratio,\
                                        first_weight_range.max() * self.channel_ratio, first_weight_range)
            last_weight_range = torch.where(last_weight_range < last_weight_range.max() * self.channel_ratio,\
                                        last_weight_range.max() * self.channel_ratio, last_weight_range)

            kernel_scale = first_weight_range.max() / (first_weight_range + eps)
            next_kernel_scale = last_weight_range.max() / (last_weight_range + eps)
            first_weight_act_range = op_act_channel_range[pair[0]]
            first_weight_act_range = torch.where(first_weight_act_range < 0.01, torch.tensor(0.01,\
                device=first_weight_act_range.device, dtype=torch.float32), first_weight_act_range)
            act_scale = first_weight_act_range.max() / (first_weight_act_range + eps)

            if algo_type == 1:
                scale = torch.min(kernel_scale, act_scale)
            elif algo_type == 2:
                kernel_scale = kernel_scale / next_kernel_scale
                act_scale = act_scale / next_kernel_scale
                scale = torch.min(kernel_scale, act_scale)
                scale = torch.min(scale, torch.tensor(ssd_min_scale, dtype=torch.float32, device=scale.device))
                scale /= scale.min()
                scale = torch.clamp(scale, 1.0, ssd_max_scale)
            else:
                kernel_scale = (kernel_scale / next_kernel_scale).sqrt()
                scale = (act_scale * kernel_scale).sqrt()
                scale = torch.clamp(scale, 1.0, ssd_max_scale)

        self.write_back(pair, scale)

    def build_observer_pair(self, pair: List[Operation]) -> Dict[Operation, OperationObserver]:
        observers = {}
        for operation in pair:
            if not isinstance(operation, QuantableOperation): continue
            observer = OperationObserver(operation=operation)
            observers[operation] = observer
        return observers

    def calibrate(
        self,
        pair: List[Operation],
        data_loader: Iterable,
        executor: BaseGraphExecutor,
        hooks: Dict[Operation, CalibrationHook],
        collate_fn: Callable,
        calib_steps: int
    ):
        calib_step = 0
        for calib_epoch in range(ceil(calib_steps / len(data_loader))):
            for _,data in enumerate(data_loader):
                if collate_fn is not None:
                    data = collate_fn(data)
                # get the input of first op
                inputs = executor.forward(data, output_names=[pair[0].inputs[0].name])
                self.run_pair(pair, inputs, hooks)
                calib_step += 1
                if calib_step >= calib_steps:
                    break

    def calibration_passive_param(self, pair: List[Operation], scale_multiplier: float=1.0):
        for op in pair:
            if not isinstance(op, QuantableOperation): continue
            if op.type in {'Conv', 'ConvTranspose', 'Gemm'}:
                if op.num_of_input == 3:
                    weight_config = op.config.input_quantization_config[1]
                    input_config = op.config.input_quantization_config[0]
                    weight_config = weight_config.dominated_by
                    input_config  = input_config.dominated_by

                    bias_config = op.config.input_quantization_config[-1]
                    if bias_config.state != QuantizationStates.PASSIVE_INIT:
                        continue

                    bias_config.scale  = weight_config.scale * input_config.scale * scale_multiplier
                    bias_config.state  = QuantizationStates.PASSIVE
                    bias_config.offset = torch.zeros_like(bias_config.scale, dtype=torch.float)
                    assert not bias_config.policy.has_property(QuantizationProperty.ASYMMETRICAL), (
                        'Negative parameter does not support ASYMMETRICAL quantization')

    def initiate_pair_state(self, pair: List[Operation]):
        for op in pair:
            if isinstance(op, QuantableOperation):
                for quant_config in op.config.input_quantization_config + op.config.output_quantization_config:
                    if quant_config.state == QuantizationStates.ACTIVATED:
                        quant_config.state = QuantizationStates.INITIAL
                    elif quant_config.state == QuantizationStates.PASSIVE:
                        quant_config.state = QuantizationStates.PASSIVE_INIT

    def dequantize_pair(self, pair: List[Operation]):
        for op in pair:
            if isinstance(op, QuantableOperation):
                op.dequantize(expire_device=None)

    def restore_quantize_state(self, pair: List[Operation]):
        for op in pair:
            if isinstance(op, QuantableOperation):
                op.restore_quantize_state(expire_device=None)

    def run_pair(self,
                 pair: List[Operation],
                 inputs: List[torch.Tensor],
                 hooks: Dict[Operation, CalibrationHook]={}) -> List[torch.Tensor]:
        for op in pair:
            inputs = inputs + [param.value for param in op.parameters]
            if isinstance(op, QuantableOperation):
                input_configs = [_ for _ in op.config.input_quantization_config]
                assert(len(inputs) == len(input_configs))
                inputs_quant = [self.quant_func(input, config) for input, config in zip(inputs, input_configs)]
                hook = hooks.get(op, None)
                if hook is not None:
                    hook.pre_forward_hook(inputs, inputs_quant, input_configs)
            else:
                inputs_quant = inputs

            f = OPERATION_FORWARD_TABLE[op.platform][op.type]
            outputs = f(op, inputs_quant)
            outputs = outputs if isinstance(outputs, (list, tuple)) else [outputs]

            if isinstance(op, QuantableOperation):
                output_configs = [_ for _ in op.config.output_quantization_config]
                outputs_quant = [self.quant_func(output, config) for output, config in zip(outputs, output_configs)]
                hook = hooks.get(op, None)
                if hook is not None:
                    hook.post_forward_hook(outputs, outputs_quant, output_configs)
                inputs = outputs_quant
            else:
                inputs = outputs
        return inputs

    # mse calculation for a list of tensors
    def calculate_mse(self, fp_res: List[torch.Tensor], quant_res: List[torch.Tensor]) -> torch.Tensor:
        losses = []
        for i in range(len(fp_res)):
            losses.append(torch_mean_square_error(fp_res[i], quant_res[i]))
        return torch.stack(losses).mean()

    @torch.no_grad()
    def test_ssd_loss(self,
                    pair: List[Operation],
                    executor: BaseGraphExecutor,
                    data_loader: Iterable,
                    collate_fn: Callable,
                    calib_steps: int
    ) -> float:
        observers = self.build_observer_pair(pair)
        hooks = {op:observers[op].hook for op in observers}
        self.calibrate(pair, data_loader, executor, hooks, collate_fn, calib_steps)
        for _, observer in observers.items():
            observer.render_quantization_config()
        pop_list = []
        for op, observer in observers.items():
            if all([type(var_observer) not in {TorchHistObserver}
                for var_observer in observer._hook._observer_table.values()]):
                    pop_list.append(op)
        for op in pop_list:
            observers.pop(op)
            hooks.pop(op)
        if len(hooks) > 0:
            self.calibrate(pair, data_loader, executor, hooks, collate_fn, calib_steps)
            for _, observer in observers.items():
                observer.render_quantization_config()
        self.calibration_passive_param(pair)
        # calculate loss
        loss = []
        for calib_epoch in range(ceil(calib_steps / len(data_loader))):
            for _,data in enumerate(data_loader):
                if collate_fn is not None:
                    data = collate_fn(data)
                # get the input of first op
                inputs = executor.forward(data, output_names=[pair[0].inputs[0].name])
                # dequant to get fp output
                self.dequantize_pair(pair)
                fp_output = self.run_pair(pair, inputs)
                # restore quant state to get quant output
                self.restore_quantize_state(pair)
                quant_output = self.run_pair(pair, inputs)
                # mse calculation
                loss.append(self.calculate_mse(fp_output, quant_output))
        return torch.stack(loss).mean().item()

    # maintain original parameter for restoration in case of a larger loss after equalization
    def collect_original_parameter(self, pair: List[Operation]) -> Dict[Variable, torch.Tensor]:
        original_weights = {}
        for op in pair:
            for var in op.inputs + op.outputs:
                if var.is_parameter:
                    original_weights[var] = var.value.clone()
        return original_weights

    # discard original fp value and store the current value as new fp value
    def store_parameter(self, pair: List[Operation]):
        for op in pair:
            if isinstance(op, QuantableOperation):
                op.store_parameter_value()

    # recover from maintained original parameters
    def recover_original_parameter(self,
                                pair: List[Operation],
                                original_weights: Dict[Variable, torch.Tensor],
    ):
        for op in pair:
            for var in op.inputs + op.outputs:
                if var.is_parameter:
                    var.value = original_weights[var]
        self.store_parameter(pair)

    def optimize(
        self,
        graph: BaseGraph,
        dataloader: Iterable,
        executor: BaseGraphExecutor,
        collate_fn: Callable,
        calib_steps: int,
        **kwargs
    ) -> None:
        # restrain maximum img number used for loss checking
        batchsize = 1
        for data in dataloader:
            if collate_fn is not None:
                data = collate_fn(data)
            if isinstance(data, torch.Tensor):
                batchsize = data.shape[0]
            elif isinstance(data, (list, tuple)):
                for value in data:
                    if isinstance(value, torch.Tensor):
                        batchsize = value.shape[0]
                        break
            elif isinstance(data, dict):
                for key, value in data.items():
                    if isinstance(value, torch.Tensor):
                        batchsize = value.shape[0]
                        break
            break

        calib_steps = min(calib_steps, ceil(200 / batchsize))

        all_pairs = self.collect_all_pairs(graph)
        if self.layer_norm:
            self.layer_weight_norm(all_pairs)

        for i in range(self.iteration):
            logger.debug(f'DFQ/SSD Equalization Iteration {i + 1}/{self.iteration}')
            for _,pair in tqdm(enumerate(all_pairs), desc=f'SSD/DFQ Equalization Iteration {i+1}/{self.iteration}', total=len(all_pairs)):
                logger.debug(f"Now Processing Pair {_ + 1}/{len(all_pairs)}: {'--'.join([op.name for op in pair])}")
                self.store_parameter(pair)

                logger.debug(f'Collecting Activation Range for Pair...')
                op_act_range = self.collect_activation_range(pair, executor, dataloader, collate_fn, calib_steps)
                logger.debug(f'Collecting Done!')

                original_weights = self.collect_original_parameter(pair)
                basic_loss = self.test_ssd_loss(pair, executor, dataloader, collate_fn, calib_steps)
                best_loss = basic_loss
                best_idx = -1

                # now apply equalization and estimate loss
                for algo in range(0,4):
                    self.one_step_equalization(pair, op_act_range, algo)
                    self.store_parameter(pair)
                    self.initiate_pair_state(pair)
                    loss = self.test_ssd_loss(pair, executor, dataloader, collate_fn, calib_steps)
                    if algo == 0:
                        logger.debug(f'DFQ Step, Loss Before Equalization {basic_loss} || Loss After Equalization {loss}')
                    else:
                        logger.debug(f'SSD Algo {algo}, Loss Before Equalization {basic_loss} || Loss After Equalization {loss}')
                    if loss < basic_loss * self.loss_threshold and loss < best_loss:
                        best_idx = algo
                        best_loss = loss
                    self.recover_original_parameter(pair, original_weights)

                if best_idx >= 0:
                    if best_idx == 0:
                        logger.debug(f'DFQ Step Activated, Loss Before Equalization {basic_loss} || Loss After Equalization {best_loss}')
                    else:
                        logger.debug(f'SSD Algo {best_idx} Activated, Loss Before Equalization {basic_loss} || Loss After Equalization {best_loss}')
                    self.one_step_equalization(pair, op_act_range, best_idx)
                    self.store_parameter(pair)
                else:
                    logger.debug('SSD and DFQ Deactivated')
                self.initiate_pair_state(pair)


from .scheduler import Step, StepDecay, Cosine, Poly # noqa F401

def epoch_to_iter(config):
    ratio = config.max_iter / config.max_epoch
    candidate_key = list(config.keys())
    for k in candidate_key:
        if k == 'lr_epochs':
            config['lr_steps'] = [round(epoch*ratio) for epoch in config[k]]
        elif k == 'warmup_epoch':
            config['warmup_steps'] = max(round(config[k]*ratio), 2)
        else:
            continue
        config.pop(k)
    config.pop('max_epoch')

def scheduler_entry(config):
    if config.type in ['StepEpoch', 'CosineEpoch']:
        config.type = config.type.replace('Epoch', '')
        epoch_to_iter(config.kwargs)
    return globals()[config.type](**config.kwargs)


# Program to buy stocks from somewhere and then sell them for another price
# and find out if that transaction was profitable. (TO THE MOON)

shares = int(input("Enter number of shares:"))
purchasePrice = float(input("Enter purchase price:"))
salePrice = float(input("Enter sale price:"))

BROKER_FEE = 0.03

# Get the total price, and then account for the brokers cost
totalAmountPaid = shares * purchasePrice
totalAmountPaid = totalAmountPaid + (totalAmountPaid * BROKER_FEE)

totalAfterSale = shares * salePrice
totalAfterSale = totalAfterSale - (totalAfterSale * BROKER_FEE)

totalReturn = totalAfterSale - totalAmountPaid

if totalAmountPaid < totalAfterSale:
    print("After the transaction, you gained {0} dollars.".format(totalReturn))
elif totalAmountPaid > totalAfterSale:
    print("After the transaction, you lost {0} dollars.".format(totalReturn * - 1))
else:
    print("You broke even.")

Chapter 02 - Input, Processing, and Output/Programming Projects/71430.py


import random
list=[]

for x in range(30):
    X = random.randint(1,15)
    list.append(x)

meu_numero = int(input("Digite um número de 1 à 15: "))
contador = 0
for x in list:
    if x == meu_numero:
        contador += 1
print(f"Eu escolhi o número: {meu_numero}. ")
print(f"Meu número repetiu, {contador} vezes. ")
print(f"Os números sorteados são: {list}")


from typing import List, Union
from .payloadObject import PayLoadObject, PayLoadStatus
from ...organizer.source import SourceObject
from gailbot.core.utils.general import (
    get_extension,
    copy)
from gailbot.core.utils.logger import makelogger
from gailbot.core.utils.media import AudioHandler
import os
from gailbot.workspace.manager import WorkspaceManager
from gailbot.configs import service_config_loader

SUPPORTED_AUDIO = service_config_loader().engines.audio_supported_format
MERGED_FILE_NAME = "merged"

logger = makelogger("audioPayload")

def load_audio_payload(
    source: SourceObject, ws_manager: WorkspaceManager) -> Union[bool, List[PayLoadObject]]:
    """ given a source object, convert it into an audio payload if the source

    Args:
        source (SourceObject): an instance of SourceObject that stores the
        datafile and setting of the transcription

    Returns:
        Union[bool, List[PayLoadObject]]: return the converted payload if the
        conversion is successful, return false other wise
    """
    if not source.setting:
        return False
    if not AudioPayload.is_supported(source.source_path()):
        return False
    try:
        return [AudioPayload(source, ws_manager)]
    except Exception as e:
        logger.error(source.__class__)
        logger.error(e, exc_info=e)
        return False

class AudioPayload(PayLoadObject):
    """
    Class for audio payload
    """
    def __init__(self, source: SourceObject, workspace: WorkspaceManager) -> None:
        super().__init__(source, workspace)

    @staticmethod
    def is_supported(file_path: str) -> bool:
        """
        Determines if a given file path has a supported file extension
        """
        logger.info(file_path)
        return get_extension(file_path) in SUPPORTED_AUDIO

    def _set_initial_status(self) -> None:
        """
        Sets the initial status of the payload object to initialized
        """
        self.status = PayLoadStatus.INITIALIZED

    def _copy_file(self) -> None:
        """
        Copies file to workspace
        """
        extension = get_extension(self.original_source)
        tgt_path =os.path.join(
            self.workspace.data_copy, f"{self.name}.{extension}")
        copy(self.original_source, tgt_path)
        self.data_files = [tgt_path]

    def _merge_audio(self):
        try:
            handler = AudioHandler()
            merged_path = handler.overlay_audios(self.data_files, self.out_dir.media_file, MERGED_FILE_NAME)
            self.merged_audio = merged_path
            assert merged_path
        except Exception as e:
            logger.error(e, exc_info=e)

    @staticmethod
    def supported_format() -> List[str]:
        """
        Contains and accesses a list of the supported formats
        """
        return SUPPORTED_AUDIO

    def __repr__(self) -> str:
        return "Audio payload"


#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Рекомендуется Python 3.x

"""
Я старался как следует прокомментировать основные моменты, чтобы в случае необходимости
в моём коде мог разобраться сотрудник, незнакомый с Python
© Илья <WST>
"""

# Импортируем важные пакеты
import time, sys, threading, requests, signal

# Время начала работы скрипта
execution_start = time.time()

# Расставим затычки-мьютексы
in_mutex = threading.Lock()
out_mutex = threading.Lock()

# Прикинемся браузером
request_headers = {
	'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Ubuntu Chromium/49.0.2623.108 Chrome/49.0.2623.108 Safari/537.36',
}

# Счётчик обработанных ссылок (для отображения прогресса)
counter = 0

# http://stackoverflow.com/questions/22492484/how-do-i-get-the-ip-address-from-a-http-request-using-the-requests-library
try:
	from requests.packages.urllib3.connectionpool import HTTPConnectionPool

	# Отключим ругань на невалидный сертификат
	from requests.packages.urllib3.exceptions import InsecureRequestWarning
	requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
except:
	print("Sadly, your version of Requests is too old.\nTry using the version from OTS repo: http://doc.mkpnet.ru/admin/deb/index.html")
	sys.exit(-1)

# Новый метод, который мы обезьянним методом воткнём вместо старого
def _make_request(self, conn, method, url, **kwargs):
	response = self._old_make_request(conn, method, url, **kwargs)
	sock = getattr(conn, 'sock', False)
	if sock:
		setattr(response, 'peer', sock.getpeername())
	else:
		setattr(response, 'peer', None)
	return response

# Осуществляем подмену
HTTPConnectionPool._old_make_request = HTTPConnectionPool._make_request
HTTPConnectionPool._make_request = _make_request

# Наш воркер
class Worker(threading.Thread):
	def __init__(self, thread_id, in_data, out_data, trace):
		threading.Thread.__init__(self),
		self.thread_id = thread_id
		self.in_data = in_data
		self.out_data = out_data
		self.timeout = 3
		self.total_count = len(in_data)
		self.trace = trace

	def select_unprocessed(self):
		with in_mutex:
			try:
				result = self.in_data.pop()
			except:
				result = None
			return result

	def report_progress(self, item):
		global counter
		counter += 1
		print(u"(%d of %d) [%s] %s" % (counter, self.total_count, item[2], item[1]))

	def process_item(self, item):
		global request_headers
		item[4] = int(time.time())

		try:
			response = requests.get(item[1], timeout = self.timeout, stream = True, headers = request_headers)
			content = response.raw.read(100000, decode_content = True)

			if b'eais.rkn.gov.ru' in content:
				item[2] = 'blocked'
			else:
				peer = response.raw._original_response.peer
				if peer is not None:
					address = peer[0]
					if address.startswith('127') or address.startswith('192.168') or address.startswith('10.10') or address == '::1' or address is None:
						item[2] = 'local-ip'
						item[3] = ''
					else:
						item[2] = 'available'
						#item[3] = content
						item[3] = ''
				else:
					item[2] = 'available'
					#item[3] = content
					item[3] = ''
		except Exception as e:
			item[2] = 'failure'

		with out_mutex:
			if self.trace:
				self.report_progress(item)
			self.out_data.append(item)

	def set_timeout(self, new_timeout):
		self.timeout = new_timeout

	def run(self):
		while True:
			item = self.select_unprocessed()
			if item is None:
				break
			else:
				self.process_item(item)

try:
	import mysql.connector as MySQLdb
except:
	print("mysql.connector is required, other drivers are not supported")
	sys.exit(-1)

# Импортируем конфигурацию
import config

# Профилирование
import resource

# Получение ID проверялки из базы данных
def get_instance_id(api = False):
	timestamp = int(time.time())
	instance_id = 0
	cursor.execute("SELECT checker_id, checker_force_scan, checker_state, checker_last_scan_time, checker_enabled, checker_threads FROM roskom_checkers WHERE REPLACE(checker_ip, '127.0.0.1', 'localhost') = (SELECT SUBSTRING_INDEX(host, ':', 1) FROM information_schema.processlist WHERE ID = connection_id()) FOR UPDATE")
	rows = cursor.fetchall()
	if len(rows) != 1:
		print("This checker instance is not registered within lanbill database. Register this checker first.")
		sys.exit(-1)
	else:
		force_scan = int(rows[0][1])
		checker_state = rows[0][2]
		checker_last_scan_time = int(rows[0][3])
		checker_enabled = rows[0][4]
		checker_threads = int(rows[0][5])

		if checker_threads != 0:
			config.THREADS = checker_threads

		if checker_enabled != 'yes':
			print("Checker is turned off in the settings section")
			sys.exit(0)

		if api:
			if force_scan == 0:
				sys.exit(0)
		else:
			if force_scan == 1:
				print("There is an unfinished scheduled scan")
				sys.exit(0)

		if checker_state == 'scanning':
			# В этом месте существует вероятность, что по факту проверка не выполняется, т.е скрипт упал или что-то такое, возможно, стоит ругаться об этом в лог, если последняя проверка была слишком давно
			# Форсируем проверку, если есть подозрение на такую ситуацию
			if checker_last_scan_time > (timestamp - 3600):
				print("Already scanning")
				sys.exit(0)

		instance_id = int(rows[0][0])
		cursor.execute("UPDATE roskom_checkers SET checker_state = 'scanning', checker_force_scan = 0, checker_last_scan_time = %s WHERE checker_id = %s", (timestamp, instance_id))

	return instance_id

# Соединение с БД
def connect_db():
	try:
		db = MySQLdb.connect(charset = 'utf8', use_unicode = True, **config.DATABASE)
		return db
	except Exception as e:
		print("MySQL connection failure\n%s" % str(e))
		print("Edit /opt/roskomcheck/config.py to configure access to lanbill database")
		sys.exit(-1)

def log_message(message, level = 'info'):
	timestamp = int(time.time())
	cursor.execute("INSERT INTO roskom_log (log_time, log_message, log_type, log_script_name) VALUES (%s, %s, %s, 'roskomcheck.py')", (timestamp, message, level))
	db.commit()

trace = True
api = False
if len(sys.argv) >= 2:
	if sys.argv[1] == 'cron':
		trace = False

	if sys.argv[1] == 'api':
		api = True

def signal_handler(signal, frame):
	timestamp = int(time.time())
	message = "instance #%d aborted by signal" % (instance_id,)
	db = connect_db()
	cursor = db.cursor()
	cursor.execute("UPDATE roskom_checkers SET checker_state = 'idle', checker_force_scan = 0 WHERE checker_id = %s", (instance_id,))
	cursor.execute("INSERT INTO roskom_log (log_time, log_message, log_type, log_script_name) VALUES (%s, %s, %s, 'roskomcheck.py')", (timestamp, message, 'info'))
	db.commit()
	print('Exitting requested')
	db.close()
	exit(0)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGQUIT, signal_handler)

# Установим соединение с БД
db = connect_db()
#db.begin()
cursor = db.cursor(buffered = True)
# Проверим, зарегистрирован ли наш экземпляр в админке. Если нет, нам незачем работать
instance_id = get_instance_id(api)
db.commit()

log_message("instance #%d started using %d threads" % (instance_id, config.THREADS))

# Получим список URL-ок, доступность которых нам требуется проверить
cursor.execute("SELECT url_text FROM roskom_url")
in_data = [[0, i[0], 'unknown', '', 0] for i in cursor.fetchall()]
out_data = []

# Можно отсоединиться от БД на время анализа
cursor.close()
db.close()

# Инициализируем наши рабочие потоки
threads = {}
for i in range(config.THREADS):
	threads[i] = Worker(i, in_data, out_data, trace)
	threads[i].set_timeout(config.HTTP_TIMEOUT)
	threads[i].setDaemon(True)

# Разветвляемся
for index, thread in threads.items():
	thread.start()

# Соединяемся
for index, thread in threads.items():
	thread.join()

# На этом этапе у нас сформирована статистика в массиве out_data, получим данные для внесения в БД
timestamp = int(time.time())
total = len(out_data)
available = [i for i in out_data if i[2] == 'available']
unavailable = [i for i in out_data if i[2] in ['blocked', 'failure', 'local-ip']]

# Предварительная оценка ресурсов для записи в лог
stat = resource.getrusage(resource.RUSAGE_SELF)

# Время окончания работы скрипта, не считая финальную транзакцию
execution_end = time.time()
execution_time = execution_end - execution_start
execution_minutes = int(execution_time / 60)
execution_seconds = (execution_time - (execution_minutes * 60))

# Установим соединение с БД для сохранения результатов
db = connect_db()
cursor = db.cursor()
cursor.execute("INSERT INTO roskom_checker_scans (scan_checker_id, scan_when, scan_urls_total, scan_urls_available, scan_urls_unavailable, scan_time) VALUES (%s, %s, %s, %s, %s, %s)", (instance_id, timestamp, total, len(available), len(unavailable), int(execution_time)))
cursor.execute("UPDATE roskom_checkers SET checker_last_scan_time = %s, checker_last_scan_id = %s, checker_state = 'idle', checker_force_scan = 0 WHERE checker_id = %s", (timestamp, cursor.lastrowid, instance_id))
cursor.execute("DELETE FROM roskom_available_links WHERE link_checker_id = %s", (instance_id,))

for link in available:
	# Внесём ссылку в список доступных при последней проверке
	cursor.execute("INSERT INTO roskom_available_links (link_url, link_body, link_checker_id) VALUES (%s, %s, %s)", (link[1], link[3], instance_id))

	# Обновим статистику в roskom_url_stat
	cursor.execute("UPDATE roskom_url_stat SET us_ptime = %s, us_checker_id = %s, us_pcount = us_pcount + 1 WHERE us_hash = MD5(%s)", (link[4], instance_id, link[1]))

log_message("instance #%d finished, taking %d kb RES and %dm %.2fs" % (instance_id, stat.ru_maxrss, execution_minutes, execution_seconds))
db.commit()
cursor.close()
db.close()

if trace:
	# Выведем статистику проверки
	print("\n Scan report")
	print(" Total: %d\n Unavailable: %d\n Available: %d\n" % (total, len(unavailable), len(available)))

if config.ENABLE_LOG:
	with open(config.LOG, "w") as handle:
		handle.write("\n".join([i[1] for i in available]))
		handle.write("\n")


import math

import numpy as np
import pandas as pd
import scipy.interpolate
import plotly.express as px
import plotly.graph_objects as go
import matplotlib.pyplot as plt
from matplotlib import animation
from colorcet import glasbey
import sklearn.datasets
import umap
import umap.plot
import umap.utils as utils
import umap.aligned_umap
from utils import load_embedded_data
from utils.embedding import SentenceTransformerBackend
from utils.topics import FrankenTopic, UMAPArgs, VectorizerArgs, TSNEArgs, KMeansArgs, HDBSCANArgs

TILE_SIZE = 5

def axis_bounds(embedding):
    left, right = embedding.T[0].min(), embedding.T[0].max()
    bottom, top = embedding.T[1].min(), embedding.T[1].max()
    adj_h, adj_v = (right - left) * 0.1, (top - bottom) * 0.1
    return [left - adj_h, right + adj_h, bottom - adj_v, top + adj_v]

def plot_series(mapper_):
    N_COLS = 3
    num_rows = math.ceil(len(mapper_.embeddings_) / N_COLS)

    fig, axs = plt.subplots(num_rows, N_COLS, figsize=(TILE_SIZE * N_COLS, TILE_SIZE * num_rows))
    ax_bound = axis_bounds(np.vstack(mapper_.embeddings_))
    for i, ax in enumerate(axs.flatten()):
        if i >= len(mapper_.embeddings_):  # skip last one if uneven number
            continue
        # current_target = ordered_target[150 * i:min(ordered_target.shape[0], 150 * i + 400)]
        len_prev = 0 if i == 0 else len(mapper_.embeddings_[i - 1])
        cmap = plt.get_cmap('tab10').colors
        ax.scatter(*mapper_.embeddings_[i].T,
                   c=[cmap[0] if j < len_prev else cmap[1] for j in range(len(mapper_.embeddings_[i]))],
                   s=2, cmap="tab10")  # c=current_target,
        ax.axis(ax_bound)
        ax.set(xticks=[], yticks=[])

    plt.tight_layout()

def plot_animation(mapper_: umap.AlignedUMAP, frames_per_slice=20):
    n_slices = len(mapper_.embeddings_)
    ax_bound = axis_bounds(np.vstack(mapper_.embeddings_))

    all_frame_layouts = []

    for slice_i in range(n_slices - 1):
        slice_j = slice_i + 1
        layout_i = mapper_.embeddings_[slice_i]
        layout_j = mapper_.embeddings_[slice_j]
        mapping = mapper_.dict_relations_[slice_i]

        # mask for all the things included in the next slice
        mask = np.array([idx in mapping for idx in range(len(layout_i))], dtype=bool)

        from_layout = np.vstack([layout_i[list(mapping.keys())],
                                 layout_i[~mask]])
        to_layout = np.vstack([layout_j[list(mapping.values())],
                               layout_i[~mask]])

        z = np.linspace(0, 1.0, frames_per_slice)
        xs = scipy.interpolate.interp1d(np.linspace(0, 1.0, 2),
                                        np.vstack([from_layout[:, 0], to_layout[:, 0]]).T)(z)
        ys = scipy.interpolate.interp1d(np.linspace(0, 1.0, 2),
                                        np.vstack([from_layout[:, 1], to_layout[:, 1]]).T)(z)

        frame_layouts = np.array([xs, ys]).T
        all_frame_layouts.append(frame_layouts)

    all_frame_layouts.append(np.array([mapper_.embeddings_[-1] for _ in range(frames_per_slice)]))

    fig = plt.figure(figsize=(TILE_SIZE, TILE_SIZE), dpi=150)
    ax = fig.add_subplot(1, 1, 1)

    scat = ax.scatter([], [], s=2)
    text = ax.text(ax_bound[0] + 0.5, ax_bound[2] + 0.5, '')
    ax.axis(ax_bound)
    ax.set(xticks=[], yticks=[])
    plt.tight_layout()

    num_frames = frames_per_slice * len(all_frame_layouts)

    def animate(i):
        slice_k = i // frames_per_slice
        slice_l = i % frames_per_slice
        scat.set_offsets(all_frame_layouts[slice_k][slice_l])
        text.set_text(f'Frame {i} ({slice_k}–{slice_l})')
        return scat

    anim = animation.FuncAnimation(
        fig,
        init_func=None,
        func=animate,
        frames=num_frames,
        interval=40)
    anim.save('data/alumap.gif', writer="pillow")
    plt.close(anim._fig)

if __name__ == '__main__':
    limit = 10000

    embedding_backend = SentenceTransformerBackend
    embedding_model = 'paraphrase-multilingual-MiniLM-L12-v2'

    tweets, embeddings = load_embedded_data(cache_dir='data/',
                                            db_file='data/identifier.sqlite',
                                            limit=limit,
                                            remove_urls=True,
                                            remove_nonals=True,
                                            remove_hashtags=True,
                                            remove_mentions=True,
                                            backend=embedding_backend,
                                            model=embedding_model)

    print('Sorting...')
    # data = list(sorted(zip(tweets.tweets, embeddings), key=lambda x: x[0].date))
    data = list(zip(tweets.tweets, embeddings))
    N_SLICES = 6
    SLICE_SIZE = math.ceil(len(data) / N_SLICES)

    print('Slicing...')
    # two modes:
    #  a) data always lives on
    #  b) data decays (only some overlap between slices)
    MODE = 'a'
    if MODE == 'a':
        slices = [data[:(i + 1) * SLICE_SIZE] for i in range(N_SLICES)]
        sliced_embeddings = [np.array([e for _, e in s]) for s in slices]
    else:
        # FIXME: no overlap yet
        slices = [data[i * SLICE_SIZE:(i + 1) * SLICE_SIZE] for i in range(N_SLICES)]
        sliced_embeddings = [np.array([e for _, e in s]) for s in slices]

    print('Building relations dict')
    relation_dicts = [{i: i for i in range(len(sliced_embeddings[si]))}
                      for si in range(N_SLICES - 1)]

    print('Running UMAP...')
    alumap = umap.AlignedUMAP(verbose=True,
                              alignment_window_size=2,
                              alignment_regularisation=0.001,
                              n_neighbors=50,
                              min_dist=0.5,
                              metric='cosine')
    mapper = alumap.fit(sliced_embeddings, relations=relation_dicts)

    print('Plotting...')
    plot_series(mapper)
    plt.show()
    plt.close()

    plot_animation(mapper, frames_per_slice=30)


from matplotlib import pyplot as plt
import numpy as np

import figlatex
import afterpulse_tile21
import textbox

lengths = {
    5.5: 128,
    7.5: 64,
    9.5: 64
}

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

figs = []

for i, (vov, length) in enumerate(lengths.items()):

    ap21 = afterpulse_tile21.AfterPulseTile21(vov)

    fig = plt.figure(num=f'figpe2-{i}', clear=True, figsize=[8, 4])

    ap21.sim.hist('mainheight', f'~saturated & ~closept & (mainpos >= 0) & (length == {length})', yscale='log', fig=fig, nbins=1000, selection=False)
    ax, = fig.get_axes()
    ax.set_xlabel('Laser peak height')
    textbox.textbox(ax, f'{vov} VoV\nFilter {length} ns', fontsize='medium', loc='center right')
    ilen = np.searchsorted(ap21.sim.filtlengths, length)
    boundaries = ap21.sim.computenpeboundaries(ilen)
    afterpulse_tile21.vlines(ax, boundaries, linestyle=':')

    figs.append(fig)

for fig in figs:
    fig.tight_layout()
    fig.show()

figlatex.save(figs)


# celery配置文件
import os
import logging.config

# 指定消息队列

broker_url = f"redis://{os.environ.get('REDIS_HOST', '192.168.239.128')}/14"

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# logging
LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'standard': {
            'format': '%(levelname)s %(asctime)s %(module)s %(lineno)d %(message)s'
        },
        'verbose': {
            'format': '%(asctime)s [%(threadName)s:%(thread)d] [%(name)s:%(lineno)d] [%(levelname)s]- %(message)s'
        },
        'simple': {
            'format': '%(levelname)s %(module)s %(lineno)d %(message)s'
        },
    },
    'handlers': {
        'file': {  # 向文件中输出日志
            'level': 'INFO',
            'class': 'logging.handlers.RotatingFileHandler',
            'filename': os.path.join(os.path.dirname(BASE_DIR), "logs/celery.log"),  # 日志文件的位置
            'maxBytes': 300 * 1024 * 1024,
            'backupCount': 10,
            'formatter': 'verbose'
        },
        'sms': {  # 向文件中输出日志
            'level': 'INFO',
            'class': 'logging.handlers.RotatingFileHandler',
            'filename': os.path.join(os.path.dirname(BASE_DIR), "logs/sms.log"),  # 日志文件的位置
            'maxBytes': 300 * 1024 * 1024,
            'backupCount': 10,
            'formatter': 'verbose'
        },
        'console': {  # 向终端中输出日志
            'level': 'INFO',
            'filters': ['require_debug_true'],
            'class': 'logging.StreamHandler',
            'formatter': 'simple'
        },
    },
    'loggers': {
        'sms': {
            'handlers': ['sms', 'console'],
            'level': 'INFO',
            'propagate': True
        },
    },
}


from django.contrib.auth import get_user_model
from rest_framework import serializers
from .models import Blog

class BlogSerializer(serializers.ModelSerializer):
    created_on = serializers.DateTimeField(format="%Y-%m-%d", read_only=True)
    content = serializers.SerializerMethodField()
    class Meta:
        model = Blog
        fields = ('id', 'title', 'img', 'text', 'content', 'created_on')

    def get_content(self, obj):
        return obj.convert_markdown_to_html()


import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...contextagg import (
    # SpatialCGNL,
    LocalAttenModule,
    CrissCrossAttention,
    # SmallLocalAttenModule,
)
from torch.nn import BatchNorm2d, BatchNorm1d
from .gcpa_gald import FAM
from ...encoders import hardnet
from ..transunet import (Transformer, Conv2dReLU)
class GCPATrans(nn.Module):
    def __init__(self, config, img_size=352):
        super(GCPATrans, self).__init__()

        self.hardnet = hardnet(arch=68)
        self.classifier = config.classifier
        self.transformer = Transformer(config, img_size,False)
        self.conv_more = Conv2dReLU(
            768, 256,
            kernel_size=3,
            padding=1,
            use_batchnorm=True,
        )
        self.config = config

        inplanes = 1024
        interplanes = 256

        self.conva = nn.Sequential(
            nn.Conv2d(inplanes, interplanes, 3, padding=1, bias=False),
            BatchNorm2d(interplanes),
            nn.ReLU(interplanes),
        )

        self.fam45 = FAM(640, interplanes, interplanes, interplanes)
        self.fam34 = FAM(320, interplanes, interplanes, interplanes)
        self.fam23 = FAM(128, interplanes, interplanes, interplanes)

        self.linear5 = nn.Conv2d(interplanes, 1, kernel_size=3, stride=1, padding=1)
        self.linear4 = nn.Conv2d(interplanes, 1, kernel_size=3, stride=1, padding=1)
        self.linear3 = nn.Conv2d(interplanes, 1, kernel_size=3, stride=1, padding=1)
        self.linear2 = nn.Conv2d(interplanes, 1, kernel_size=3, stride=1, padding=1)

        self.local_attention_4 = LocalAttenModule(interplanes)
        self.local_attention_3 = LocalAttenModule(interplanes)
        self.local_attention_2 = LocalAttenModule(interplanes)

    def forward(self, x):
        hardnetout = self.hardnet(x)
        transout, attn_weights, features = self.transformer(x)

        out2 = hardnetout[0]  # [24, 128, 88, 88]
        out3 = hardnetout[1]  # [24, 320, 44, 44]
        out4 = hardnetout[2]  # [24, 640, 22, 22]
        out5 = hardnetout[3]  # [24, 1024, 11, 11]
        out5 = self.conva(out5)

        z = transout
        print(z.shape)
        B, n_patch, hidden = z.size()
        h, w = int(np.sqrt(n_patch)), int(np.sqrt(n_patch))
        z = z.permute(0, 2, 1)
        z = z.contiguous().view(B, hidden, h, w)
        z = self.conv_more(z)  # bs, 256, 11, 11
        z = F.interpolate(z, scale_factor=0.5, mode='bilinear')

        # LD
        out4_c = self.local_attention_4(z)  # bs, 256, 11, 11

        out3_c = self.local_attention_3(z)  # bs, 256, 11, 11

        out2_c = self.local_attention_2(z)  # bs, 256, 11, 11

        # out
        out4 = self.fam45(out4, out5, out4_c)
        out3 = self.fam34(out3, out4, out3_c)
        out2 = self.fam23(out2, out3, out2_c)
        # we use bilinear interpolation instead of transpose convolution
        out5 = F.interpolate(self.linear5(out5), size=x.size()[2:], mode="bilinear")
        out4 = F.interpolate(self.linear4(out4), size=x.size()[2:], mode="bilinear")
        out3 = F.interpolate(self.linear3(out3), size=x.size()[2:], mode="bilinear")
        out2 = F.interpolate(self.linear2(out2), size=x.size()[2:], mode="bilinear")
        return out5, out4, out3, out2


# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from keras import initializers, regularizers
from keras.engine.topology import InputSpec
from keras import backend as K

import numpy as np

import warnings

from ..capsule import Module
from ..utils.caps_utils import mixed_shape
from ..utils import dissimilarity_funcs as diss_funcs
from LVQ import constraints
from ..utils import pre_training

class PointPrototype(Module):
    def __init__(self,
                 prototype_initializer='TruncatedNormal',
                 prototype_regularizer=None,
                 prototype_constraint=None,  # That's similar to a normalization.
                 signal_output='signals',
                 linear_factor=0.5,  # linear factor  (None --> always add) prev_diss*(1-alpha) + alpha*curr_diss
                 signal_regularizer=None,
                 diss_regularizer=None,
                 **kwargs):

        if linear_factor is None or (isinstance(linear_factor, float) and 0 < linear_factor < 1):
            self.linear_factor = linear_factor
        else:
            raise TypeError("The linear factor for the dissimilarity value must be None or a float in the range of "
                            "(0,1). You provide: " + str(linear_factor))

        self.prototype_initializer = initializers.get(prototype_initializer)
        self.prototype_regularizer = regularizers.get(prototype_regularizer)
        self.prototype_constraint = constraints.get(prototype_constraint)

        self.prototypes = None

        self.signal_output = signal_output

        self.output_regularizers = [regularizers.get(signal_regularizer),
                                    regularizers.get(diss_regularizer)]

        # be sure to call this at the end
        super(PointPrototype, self).__init__(module_input=True,
                                             module_output=True,
                                             support_sparse_signal=True,
                                             support_full_signal=True,
                                             **self._del_module_args(**kwargs))

    def _build(self, input_shape):
        if not self.built:
            if input_shape[0][1] != self.capsule_number:
                raise ValueError('The capsule number provided by input_shape is not equal the self.capsule_number: '
                                 'input_shape[0][1]=' + str(input_shape[0][1]) + ' != ' +
                                 'self.capsule_number=' + str(self.capsule_number))
            if input_shape[1][1] != self.proto_number:
                raise ValueError('The prototype number provided by input_shape is not equal the self.proto_number: '
                                 'input_shape[1][1]=' + str(input_shape[1][1]) + ' != ' +
                                 'self.proto_number=' + str(self.proto_number))
            # the signal dimension is input_shape[0][3:]
            self.prototypes = self.add_weight(shape=(self.proto_number,) + tuple(input_shape[0][3:]),
                                              initializer=self.prototype_initializer,
                                              regularizer=self.prototype_regularizer,
                                              constraint=self.prototype_constraint,
                                              name='prototypes')

            self.input_spec = [InputSpec(shape=(None,) + tuple(input_shape[0][1:])),
                               InputSpec(shape=(None,) + tuple(input_shape[1][1:]))]

    def _build_sparse(self, input_shape):
        if self.signal_output != 'signals':
            raise ValueError('If the signal is sparse, the signal_output must be `signals`.')

        # manipulate input_shape to call the full build method instead of a new implementation
        signal_shape = list(input_shape[0])
        signal_shape[1] = self.capsule_number

        self._build([tuple(signal_shape), input_shape[1]])

        self.input_spec = [InputSpec(shape=(None,) + tuple(input_shape[0][1:])),
                           InputSpec(shape=(None,) + tuple(input_shape[1][1:]))]

    def _call(self, inputs, **kwargs):
        # signals_shape input: batch x capsule_number x channels x dim1 x dim2 x ... x dimN
        # after measuring the signal is (output): batch x proto_number x channels x dim1 x dim2 x ... x dimN
        signals = inputs[0]
        diss0 = inputs[1]

        if self.sparse_signal:
            signals, diss1 = self.distance_sparse(signals)
        else:
            signals, diss1 = self.distance(signals)

        with K.name_scope('dissimilarity_update'):
            if self.linear_factor is None:
                diss = diss0 + diss1
            else:
                diss = (1 - self.linear_factor) * diss0 + self.linear_factor * diss1

        return {0: signals, 1: diss}

    def _call_sparse(self, inputs, **kwargs):
        if self.signal_output != 'signals':
            raise ValueError('If the signal is sparse, the signal_output must be `signals`.')
        return self._call(inputs, **kwargs)

    def distance(self, signals):
        raise NotImplementedError

    def distance_sparse(self, signals):
        raise NotImplementedError

    def _compute_output_shape(self, input_shape):
        signal_shape = list(input_shape[0])
        diss_shape = list(input_shape[1])

        signal_shape[1] = self.proto_number

        return [tuple(signal_shape), tuple(diss_shape)]

    def _compute_output_shape_sparse(self, input_shape):
        if self.signal_output != 'signals':
            raise ValueError('If the signal is sparse, the signal_output must be `signals`.')
        return input_shape

    def get_config(self):
        config = {'prototype_initializer': initializers.serialize(self.prototype_initializer),
                  'prototype_regularizer': regularizers.serialize(self.prototype_regularizer),
                  'prototype_constraint': constraints.serialize(self.prototype_constraint),
                  'signal_output': self.signal_output,
                  'linear_factor': self.linear_factor,
                  'signal_regularizer': regularizers.serialize(self.output_regularizers[0]),
                  'diss_regularizer': regularizers.serialize(self.output_regularizers[1])}
        super_config = super(PointPrototype, self).get_config()
        return dict(list(super_config.items()) + list(config.items()))

    def pre_training(self,
                     x_train,
                     y_train=None,  # class labels [0,..,n], if given assume class correspondence and init respectively
                     capsule_inputs_are_equal=False,  # use just the first capsule channels; ignored if sparse
                     batch_version=False,
                     **kmeans_params):  # params of kMeans or batch kMeans of sklearn

        # overwrite flag if sparse
        if self.sparse_signal:
            capsule_inputs_are_equal = True

        try:
            self._is_callable()
        except ValueError:
            raise ValueError("The module " + self.name + " is not proper initialized for pre-training. Be sure that you"
                             " assigned the module to a capsule and built it by calling the capsule.")

        if self.input_spec[0].shape[1:] != x_train.shape[1:]:
            raise ValueError("The shape of x_train must be equal to the assumed input shape. You provide: "
                             + str(x_train.shape) + "; we expect: " + str(self.input_spec[0].shape))

        if y_train is not None:
            # convert y_train to non categorical
            if y_train.ndim != 1:
                y_train = np.argmax(y_train, 1)

            # check if all classes are present
            if list(np.unique(y_train)) != list(range(self.capsule_number)):
                raise ValueError("We detected that not each class is represented in the training data. Be sure that "
                                 "for each class, at least one sample is given.")

            if y_train.shape[0] != x_train.shape[0]:
                raise ValueError("The number of samples in y_train and x_train must be equal. You provide: y_train="
                                 + str(y_train.shape[0]) + " is not equal x_train=" + str(x_train.shape[0]) + ".")

        centers, _, _ = _pre_train_prototypes(x_train,
                                              y_train,
                                              capsule_inputs_are_equal,
                                              batch_version,
                                              self._proto_distrib,
                                              **kmeans_params)

        # we apply constraint here, because pre-training is like an optimization step
        if self.prototype_constraint is not None:
            centers = K.eval(self.prototype_constraint(K.variable(centers)))

        self.set_weights([centers] + self.get_weights()[1:])

# definition of prototype which consist of a data point and a matrix
class PointMatrixPrototype(PointPrototype):
    def __init__(self,
                 prototype_initializer='TruncatedNormal',
                 prototype_regularizer=None,
                 prototype_constraint=None,  # That's similar to a normalization.
                 matrix_initializer='TruncatedNormal',
                 matrix_regularizer=None,
                 matrix_constraint=None,  # That's similar to a normalization.
                 matrix_scope='local',  # local, global or capsule_wise
                 projected_atom_shape=None,
                 signal_output='signals',
                 linear_factor=0.5,  # linear factor  (None --> always add) prev_diss*(1-alpha) + alpha*curr_diss
                 signal_regularizer=None,
                 diss_regularizer=None,
                 **kwargs):

        super(PointMatrixPrototype, self).__init__(prototype_initializer=prototype_initializer,
                                                   prototype_regularizer=prototype_regularizer,
                                                   prototype_constraint=prototype_constraint,
                                                   signal_output=signal_output,
                                                   linear_factor=linear_factor,
                                                   signal_regularizer=signal_regularizer,
                                                   diss_regularizer=diss_regularizer,
                                                   **kwargs)

        # scope: local, global, capsule_wise (defines the apply scope of the transformation)
        if matrix_scope in ('local', 'global', 'capsule_wise'):
            self.matrix_scope = matrix_scope
        else:
            raise ValueError("scope must be 'local', 'global' or 'capsule_wise': You provide: " + str(matrix_scope))

        if isinstance(projected_atom_shape, (list, tuple)):
            self._projected_atom_shape = tuple(projected_atom_shape)
        elif isinstance(projected_atom_shape, int):
            self._projected_atom_shape = (projected_atom_shape,)
        elif projected_atom_shape is None:
            self._projected_atom_shape = projected_atom_shape
        else:
            raise ValueError("projected_atom_shape must be list, tuple, int or None. You provide: " +
                             str(projected_atom_shape))
        self.projected_atom_shape = None

        self.matrix_initializer = initializers.get(matrix_initializer)
        self.matrix_regularizer = regularizers.get(matrix_regularizer)
        self.matrix_constraint = constraints.get(matrix_constraint)

        self.matrices = None

        self.input_output_shape_equal = ('signals', 'protos')

    def _build(self, input_shape):
        if not self.built:
            # compute shape of matrix_shape
            if self.matrix_scope == 'local':
                self._num_maps = (self.proto_number,)
            elif self.matrix_scope == 'global':
                self._num_maps = ()
            else:  # capsule_wise
                self._num_maps = (self.capsule_number,)

            super(PointMatrixPrototype, self)._build(input_shape)

            if self._projected_atom_shape is None:
                self.projected_atom_shape = self._compute_output_shape(input_shape)[0][3:]
            else:
                self.projected_atom_shape = self._projected_atom_shape

            matrix_shape = self._num_maps + (np.prod(np.array(input_shape[0][3:], dtype=int)),
                                             np.prod(np.array(self.projected_atom_shape, dtype=int)))

            # transposed matrix (we compute x^T * A^T instead of A * x, because the signal is always given in the form
            # x^T)
            self.matrices = self.add_weight(shape=matrix_shape,
                                            initializer=self.matrix_initializer,
                                            regularizer=self.matrix_regularizer,
                                            constraint=self.matrix_constraint,
                                            name='matrices')

    def _compute_output_shape(self, input_shape):
        signal_shape = list(input_shape[0])
        diss_shape = input_shape[1]

        signal_shape[1] = self.proto_number
        if self.signal_output in self.input_output_shape_equal:
            return [signal_shape, diss_shape]
        else:
            if not self.built:
                if self._projected_atom_shape is None:
                    projected_atom_shape = tuple(input_shape[0][3:])
                else:
                    projected_atom_shape = self._projected_atom_shape
                return [tuple(signal_shape[0:3]) + projected_atom_shape, diss_shape]
            else:
                return [tuple(signal_shape[0:3]) + self.projected_atom_shape, diss_shape]

    def get_config(self):
        config = {'matrix_initializer': initializers.serialize(self.matrix_initializer),
                  'matrix_regularizer': regularizers.serialize(self.matrix_regularizer),
                  'matrix_constraint': constraints.serialize(self.matrix_constraint),
                  'matrix_scope': self.matrix_scope,
                  'projected_atom_shape': self.projected_atom_shape}
        super_config = super(PointMatrixPrototype, self).get_config()
        return dict(list(super_config.items()) + list(config.items()))

    def pre_training(self,
                     x_train,
                     y_train=None,  # class labels [0,..,n], if given assume class correspondence and init respectively
                     capsule_inputs_are_equal=False,  # use just the first capsule channels
                     batch_version=False,
                     **sk_params):
        # sk_params: {'kmeans_params': dict,
        #             'svd_params': dict}
        # kmeans_params ... params dict of kmeans_params or kmeans_params_batch; Note: in the case of a memory error the
        # method switches automatically to the batch version. In this case be sure that all parameters are accept by the
        # batch version.
        # svd_params ... params of svd_params
        # Todo: test that svd works correct (full rank and sparse rank)

        # overwrite flag if sparse
        if self.sparse_signal:
            capsule_inputs_are_equal = True

        if hasattr(sk_params, 'kmeans_params'):
            kmeans_params = sk_params['kmeans_params']
            if not isinstance(kmeans_params, dict):
                raise TypeError("The type of kmeans_params parameters in sk_params must be dict.")
            del sk_params['kmeans_params']
        else:
            kmeans_params = {}

        if hasattr(sk_params, 'svd_params'):
            svd_params = sk_params['svd_params']
            if not isinstance(svd_params, dict):
                raise TypeError("The type of the svd_params parameters in sk_params must be dict.")
            del sk_params['svd_params']
        else:
            svd_params = {}

        try:
            self._is_callable()
        except ValueError:
            raise ValueError("The module " + self.name + " is not proper initialized for pre-training. Be sure that you"
                             " assigned the module to a capsule and built it by calling the capsule.")

        if self.input_spec[0].shape[1:] != x_train.shape[1:]:
            raise ValueError("The shape of x_train must be equal to the assumed input shape. You provide: "
                             + str(x_train.shape) + "; we expect: " + str(self.input_spec[0].shape))

        if y_train is not None:
            # convert y_train to non categorical
            if y_train.ndim != 1:
                y_train = np.argmax(y_train, 1)

            # check if all classes are present
            if list(np.unique(y_train)) != list(range(self.capsule_number)):
                raise ValueError("We detected that not each class is represented in the training data. Be sure that "
                                 "for each class, at least one sample is given.")

            if y_train.shape[0] != x_train.shape[0]:
                raise ValueError("The number of samples in y_train and x_train must be equal. You provide: y_train="
                                 + str(y_train.shape[0]) + " is not equal x_train=" + str(x_train.shape[0]) + ".")

        centers, labels, clusters = _pre_train_prototypes(x_train,
                                                          y_train,
                                                          capsule_inputs_are_equal,
                                                          batch_version,
                                                          self._proto_distrib,
                                                          **kmeans_params)

        # we apply constraint here, because pre-training is like an optimization step
        if self.prototype_constraint is not None:
            centers = K.eval(self.prototype_constraint(K.variable(centers)))

        # transform clusters in accordance to the matrix scope
        if self.matrix_scope == 'global':
            clusters = [np.reshape(np.concatenate(clusters, 0), [-1, np.prod(self.input_spec[0].shape[3:], dtype=int)])]
        elif self.matrix_scope == 'local':
            clusters_ = []
            for i in range(len(clusters)):
                for l in range(max(labels[i])+1):
                    cluster_flatten = np.reshape(clusters[i], [-1, np.prod(self.input_spec[0].shape[3:], dtype=int)])
                    clusters_.append(cluster_flatten[labels[i] == l, :])
            clusters = clusters_
        else:
            clusters_ = []
            for i in range(len(clusters)):
                clusters_.append(np.reshape(clusters[i], [-1, np.prod(self.input_spec[0].shape[3:], dtype=int)]))
            clusters = clusters_

        # center the clusters
        for i in range(len(clusters)):
            clusters[i] = clusters[i] - np.mean(clusters[i], axis=0)

        if K.int_shape(self.matrices)[-1] > K.int_shape(self.matrices)[-2]:
            warnings.warn("Can't pre-train matrices if the projection shape (" + str(K.int_shape(self.matrices)[-1]) +
                          ") is higher than the input shape (" + str(K.int_shape(self.matrices)[-2]) + "). Skip " +
                          "pre-training of matrices and use the current parameters.")
            self.set_weights([centers, self.get_weights()[1]])
        else:
            n_components = np.prod(self.projected_atom_shape, dtype=int)
            matrices, valid = pre_training.svd(clusters, n_components=n_components, **svd_params)

            old_matrices = self.get_weights()[1]
            warning_sent = False
            for i in range(len(matrices)):
                if not valid[i]:
                    if not warning_sent:
                        warnings.warn("Some of the matrices were not be pre-tarined successfully. The old matrix "
                                      "instantiation is reused and the pre-trained matrix is skipped. One reason for "
                                      "that could be that your cluster doesn't contain enough points regarding your "
                                      "matrix dimension. If you are not pre-training on the whole dataset try to "
                                      "increase the number of the pre-training batch.")
                        warning_sent = True
                    matrices[i] = old_matrices[i]

            if self.matrix_scope == 'global':
                matrices = matrices[0]
            else:
                matrices = np.stack(matrices, axis=0)

            # we apply constraint here, because pre-training is like an optimization step
            if self.matrix_constraint is not None:
                matrices = K.eval(self.matrix_constraint(K.variable(matrices)))

            self.set_weights([centers, matrices] + self.get_weights()[2:])

class MinkowskiDistance(PointPrototype):
    def __init__(self,
                 order_p=2,  # could be np.inf
                 prototype_initializer='TruncatedNormal',
                 prototype_regularizer=None,
                 prototype_constraint=None,  # That's similar to a normalization.
                 signal_output='signals',  # what should be send to the output (signals, protos, params)
                 epsilon=K.epsilon(),  # used for stabilization of the sqrt
                 squared_dissimilarity=False,  # output the squared distance (without reciprocal power)
                 linear_factor=0.5,  # linear factor  (None --> always add) prev_diss*(1-alpha) + alpha*curr_diss
                 signal_regularizer=None,
                 diss_regularizer=None,
                 **kwargs):

        valid_signal_output = ('signals', 'protos')
        if signal_output not in valid_signal_output:
            raise ValueError("signal_output must be in " + str(valid_signal_output) + ". You provide: "
                             + str(signal_output))

        if isinstance(order_p, (int, float)) or order_p == np.inf:
            if order_p > 0:
                self.order_p = order_p
            else:
                raise ValueError("The order p of the Minkowski distance must be greater than 0. You provide: "
                                 + str(order_p))
        else:
            raise TypeError("order_p must be int float or numpy.inf. You provide: " + str(order_p))

        self.squared_dissimilarity = squared_dissimilarity
        self.epsilon = epsilon

        super(MinkowskiDistance, self).__init__(prototype_initializer=prototype_initializer,
                                                prototype_regularizer=prototype_regularizer,
                                                prototype_constraint=prototype_constraint,
                                                signal_output=signal_output,
                                                linear_factor=linear_factor,
                                                signal_regularizer=signal_regularizer,
                                                diss_regularizer=diss_regularizer,
                                                **kwargs)

    def distance(self, signals):
        # Todo: make test with and without sqrt by this simple maximum stabilization if it works think about more
        # complex implementations for stabilization (test convergence behavior with sqrt and w/o)

        # tile capsule regarding proto distribution
        if self.proto_number != self.capsule_number:
            with K.name_scope('signal_preprocessing'):
                # vector_shape for permute commands
                vector_shape = list(range(3, self.input_spec[0].ndim))
                signals = K.gather(K.permute_dimensions(signals, [1, 0, 2] + vector_shape),
                                   self._capsule_extension)
                signals = K.permute_dimensions(signals, [1, 0, 2] + vector_shape)

        diss = diss_funcs.minkowski_distance(signals=signals,
                                             protos=self.prototypes,
                                             order_p=self.order_p,
                                             squared=self.squared_dissimilarity,
                                             epsilon=self.epsilon)

        with K.name_scope('get_signals'):
            if self.signal_output == 'protos':
                signal_shape = mixed_shape(signals)
                signals = K.tile(K.expand_dims(K.expand_dims(self.prototypes, 1), 0),
                                 [K.shape(signals)[0], 1, signal_shape[2]]
                                 + list(np.ones((self.input_spec[0].ndim - 3,), dtype=int)))
            else:
                signals = signals

        return signals, diss

    def distance_sparse(self, signals):
        diss = diss_funcs.minkowski_distance(signals=signals,
                                             protos=self.prototypes,
                                             order_p=self.order_p,
                                             squared=self.squared_dissimilarity,
                                             epsilon=self.epsilon)

        return signals, diss

    def get_config(self):
        config = {'order_p': self.order_p,
                  'squared_dissimilarity': self.squared_dissimilarity,
                  'epsilon': self.epsilon}
        super_config = super(MinkowskiDistance, self).get_config()
        return dict(list(super_config.items()) + list(config.items()))

class TangentDistance(PointMatrixPrototype):
    def __init__(self,
                 projected_atom_shape,
                 prototype_initializer='TruncatedNormal',
                 prototype_regularizer=None,
                 prototype_constraint=None,
                 matrix_initializer='TruncatedNormal',
                 matrix_regularizer=None,
                 matrix_scope='local',  # local, global or capsule_wise
                 signal_output='signals',  # what should be send to the output (signals, protos, projected_signals, projected_protos )
                 epsilon=K.epsilon(),
                 squared_dissimilarity=False,
                 linear_factor=0.5,  # linear factor  (None --> always add) prev_diss*(1-alpha) + alpha*curr_diss
                 signal_regularizer=None,
                 diss_regularizer=None,
                 **kwargs):

        valid_signal_output = ('signals', 'projected_signals', 'parameterized_signals')
        if signal_output not in valid_signal_output:
            raise ValueError("signal_output must be in " + str(valid_signal_output) + ". You provide: "
                             + str(signal_output))

        self.squared_dissimilarity = squared_dissimilarity
        self.epsilon = epsilon

        if projected_atom_shape is None:
            raise ValueError("projected_atom_shape must be unequal None.")

        if 'matrix_constraint' in kwargs:
            del kwargs['matrix_constraint']

        super(TangentDistance, self).__init__(prototype_initializer=prototype_initializer,
                                              prototype_regularizer=prototype_regularizer,
                                              prototype_constraint=prototype_constraint,
                                              matrix_initializer=matrix_initializer,
                                              matrix_regularizer=matrix_regularizer,
                                              matrix_constraint='orthogonalization',
                                              matrix_scope=matrix_scope,
                                              projected_atom_shape=projected_atom_shape,
                                              signal_output=signal_output,
                                              linear_factor=linear_factor,
                                              signal_regularizer=signal_regularizer,
                                              diss_regularizer=diss_regularizer,
                                              **kwargs)

        self.input_output_shape_equal = ('signals', 'projected_signals')

    def _build(self, input_shape):
        super(TangentDistance, self)._build(input_shape)

        if np.prod(input_shape[0][3:]) <= np.prod(self.projected_atom_shape):
            raise ValueError("The dimension of the projected shape must be lower than input_shape. You "
                             "provide: np.prod(signal_shape[3:])=" + str(np.prod(input_shape[0][3:])) + " < " +
                             "np.prod(self.projected_atom_shape)=" + str(np.prod(self.projected_atom_shape)))

    def distance(self, signals):
        # tile capsule regarding proto distribution
        if self.proto_number != self.capsule_number:
            with K.name_scope('signal_preprocessing'):
                # vector_shape for permute commands
                vector_axes = list(range(3, self.input_spec[0].ndim))
                signals = K.gather(K.permute_dimensions(signals, [1, 0, 2] + vector_axes),
                                   self._capsule_extension)
                signals = K.permute_dimensions(signals, [1, 0, 2] + vector_axes)

        if self.matrix_scope == 'capsule_wise':
            with K.name_scope('subspace_preprocessing'):
                matrices = K.gather(self.matrices, self._capsule_extension)
        else:
            matrices = self.matrices

        diss = diss_funcs.tangent_distance(signals=signals,
                                           protos=self.prototypes,
                                           subspaces=matrices,
                                           squared=self.squared_dissimilarity,
                                           epsilon=self.epsilon)

        with K.name_scope('get_signals'):
            if self.signal_output == 'signals':
                signals = signals

            elif self.signal_output == 'projected_signals':
                signals = self._get_projected_signals(signals, matrices)

            else:  # 'parameterized_signals'
                signals = self._get_parameterized_signals(signals, matrices)

        return signals, diss

    def distance_sparse(self, signals):
        if self.matrix_scope == 'capsule_wise':
            with K.name_scope('subspace_preprocessing'):
                matrices = K.gather(self.matrices, self._capsule_extension)
        else:
            matrices = self.matrices

        diss = diss_funcs.tangent_distance(signals=signals,
                                           protos=self.prototypes,
                                           subspaces=matrices,
                                           squared=self.squared_dissimilarity,
                                           epsilon=self.epsilon)

        return signals, diss

    def get_projected_signals(self, signals):
        if self.matrix_scope == 'capsule_wise':
            with K.name_scope('matrix_preprocessing'):
                matrices = K.gather(self.matrices, self._capsule_extension)
        else:
            matrices = self.matrices

        return self._get_projected_signals(signals, matrices)

    # affine_projection: w + W * W.T * (v - w)
    def _get_projected_signals(self, signals, matrices):
        # signal must be of shape (batch x protos x channels x dim1 x ... x dimN
        with K.name_scope('get_projected_signals'):
            signal_shape = mixed_shape(signals)
            atom_axes = list(range(3, len(signal_shape)))

            signals = K.permute_dimensions(signals, [0, 2, 1] + atom_axes)
            diff = signals - self.prototypes

            if self.matrix_scope == 'global':
                with K.name_scope('projector'):
                    projector = K.dot(matrices, K.transpose(matrices))

                with K.name_scope('tangentspace_projections'):
                    diff = K.reshape(diff, (signal_shape[0] * signal_shape[2], signal_shape[1], -1))
                    projected_diff = K.dot(diff, projector)
                    projected_diff = K.reshape(projected_diff,
                                               (signal_shape[0], signal_shape[2], signal_shape[1]) + signal_shape[3:])

                    matching_protos = self.prototypes + projected_diff
                    matching_protos = K.permute_dimensions(matching_protos, [0, 2, 1] + atom_axes)

            else:  # local or capsule_wise
                with K.name_scope('projector'):
                    projector = K.batch_dot(matrices, matrices, [2, 2])

                with K.name_scope('tangentspace_projections'):
                    diff = K.reshape(diff, (signal_shape[0] * signal_shape[2], signal_shape[1], -1))
                    diff = K.permute_dimensions(diff, [1, 0, 2])
                    projected_diff = K.batch_dot(diff, projector)
                    projected_diff = K.permute_dimensions(projected_diff, [1, 0, 2])
                    projected_diff = K.reshape(projected_diff,
                                               (signal_shape[0], signal_shape[2], signal_shape[1]) + signal_shape[3:])

                    matching_protos = self.prototypes + projected_diff
                    matching_protos = K.permute_dimensions(matching_protos, [0, 2, 1] + atom_axes)

            return matching_protos

    def get_parameterized_signals(self, signals):
        if self.matrix_scope == 'capsule_wise':
            with K.name_scope('matrix_preprocessing'):
                matrices = K.gather(self.matrices, self._capsule_extension)
        else:
            matrices = self.matrices

        return self._get_parameterized_signals(signals, matrices)

    # params: W.T * (v - w)
    def _get_parameterized_signals(self, signals, matrices):
        # signal must be of shape (batch x protos x channels x dim1 x ... x dimN
        with K.name_scope('get_projected_signals'):
            signal_shape = mixed_shape(signals)
            atom_axes = list(range(3, len(signal_shape)))

            signals = K.permute_dimensions(signals, [0, 2, 1] + atom_axes)
            diff = signals - self.prototypes

            if self.matrix_scope == 'global':
                with K.name_scope('tangentspace_parameters'):
                    diff = K.reshape(diff, (signal_shape[0] * signal_shape[2], signal_shape[1], -1))
                    params = K.dot(diff, matrices)

                    params = K.reshape(params,
                                       (signal_shape[0], signal_shape[2], signal_shape[1]) + self.projected_atom_shape)
                    params = K.permute_dimensions(params, [0, 2, 1] + atom_axes)

            else:  # local or capsule_wise
                with K.name_scope('tangentspace_parameters'):
                    diff = K.reshape(diff, (signal_shape[0] * signal_shape[2], signal_shape[1], -1))
                    diff = K.permute_dimensions(diff, [1, 0, 2])
                    params = K.batch_dot(diff, matrices)

                    params = K.reshape(params,
                                       (signal_shape[1], signal_shape[0], signal_shape[2]) + self.projected_atom_shape)
                    params = K.permute_dimensions(params, [1, 0, 2] + atom_axes)

            return params

    def get_config(self):
        config = {'squared_dissimilarity': self.squared_dissimilarity,
                  'epsilon': self.epsilon}
        super_config = super(TangentDistance, self).get_config()
        return dict(list(super_config.items()) + list(config.items()))

def _pre_train_prototypes(x_train,
                          y_train,  # non-categorical
                          capsule_inputs_are_equal,
                          batch_version,
                          proto_distrib,
                          **kmeans_params):
    # signals_shape: samples x capsule_number x channels x dim1 x dim2 x ... x dimN

    signal_shape = x_train.shape[3:]
    proto_number = sum([len(x) for x in proto_distrib])

    x_trains = []

    # all the cases for k-means computation
    if y_train is None and capsule_inputs_are_equal:
        # samples x dim1 x dim2 x ... x dimN
        x = x_train[:, 0, :]
        n_clusters = proto_number
        x = np.reshape(x, (-1, np.prod(signal_shape)))
        x_trains.append(x)
    else:
        n_clusters = []

        for i, p in enumerate(proto_distrib):
            n_clusters.append(len(p))

            if y_train is None and not capsule_inputs_are_equal:
                # capsule_number x samples x dim1 x dim2 x ... x dimN
                x = x_train[:, i, :]

            elif y_train is not None and capsule_inputs_are_equal:
                # samples x channels x dim1 x dim2 x ... x dimN
                x = x_train[y_train == i, 0, :]

            else:
                # samples x capsule_number x channels x dim1 x dim2 x ... x dimN
                x = x_train[y_train == i, i, :]

            x = np.reshape(x, (-1, np.prod(signal_shape)))
            x_trains.append(x)

    centers, labels = pre_training.kmeans(x_trains, n_clusters, batch_version, **kmeans_params)

    centers = np.concatenate(centers, 0)
    # reshape back to real shape
    centers = np.reshape(centers, (proto_number, ) + signal_shape)

    return centers, labels, x_trains


# ドイツの言語クラスをインポート
from spacy.lang.de import German

# nlpオブジェクトを作成
nlp = German()

# テキストを処理（ドイツ語で「よろしく！」の意味）
doc = nlp("Liebe Grüße!")

# docのテキストをプリント
print(doc.text)


import gym
import numpy as np
import random
import matplotlib.pyplot as plt

env = gym.make('MountainCar-v0')

# Initialize q-table values to 0
discretization = 15
pos_discretization = 15
Q = np.zeros((pos_discretization,discretization, env.action_space.n))    # (states_space, action_space)

# Initialize hyperparameters
gamma = 0.98         # Discount factor used to balance immediate and future reward.
lr = 0.1            # learning rate
epsilon = 0.9         # rate of exploration where 1 is completely random and 0 is deterministic
epsilon_decay = 0.02
# Initialize list for plotting running mean
rewards = []
N = 100     # Number of episodes for the running mean

def running_mean(x, N):
    """function for creating list with running mean"""
    cumsum = np.cumsum(np.insert(x, 0, 0))
    return (cumsum[N:] - cumsum[:-N]) / float(N)

def get_state(cart_position, tip_speed):
    """Returns index of the given speed state in the Q-table"""
    positions = np.linspace(-1.2,0.6,num=pos_discretization)
    states = np.linspace(-0.07,0.07,num=discretization)   # Discretized table of speeds.
    for rounded_cart_position in positions:
        if rounded_cart_position >= cart_position:
            cart_index = np.where(positions == rounded_cart_position)[0][0]
            for rounded_tip_speed in states:
                if rounded_tip_speed >= tip_speed:
                    index = np.where(states == rounded_tip_speed)
                    index = index[0][0]
                    return cart_index, index
    print("NONE",cart_position,cart_index)
    print(str(index))
    return None

def better_action(cart_position, tip_speed):
    """Returns index of best action according to Q-table for a given tip speed"""
    possible_actions_in_state = Q[get_state(cart_position, tip_speed)]
    action_of_choice = np.argmax(possible_actions_in_state)
    return action_of_choice

def update_q(state, action, new_state, cart_pos, new_cart_pos, reward=1):
    Q[cart_pos, state, action] = Q[cart_pos, state, action] + lr * (reward + gamma * np.max(Q[new_cart_pos, new_state, :]) - Q[cart_pos, state, action])

for i_episode in range(5000):
    observation = env.reset()
    episode_reward = 0

    if 4900 < i_episode :
        # Drop exploration
        epsilon = 0

    if i_episode%100 == 0 and i_episode != 0:
        #Give some feedback for improvement of the algorithm
        print(i_episode)
        print("best run was", max(rewards[-99:]))

    for timestep in range(500):
        if i_episode > 4900:
            env.render()
            pass

        #if len(rewards) > 1:

            #if episode_reward >= max(rewards) and episode_reward > -150:
                #Implementing dynamic epsilon so that exploration decays over time.
             #   epsilon -= 0.006

        tip_speed = observation[1]
        cart_position = observation[0]

        if random.uniform(0, 1) < epsilon:      # Chooses random action often if epsilon is high
            action = env.action_space.sample()

        else: #Chooses action based on previous exploration
            action = better_action(cart_position, tip_speed)

        # Gather information about changes in environment
        cart_state, state = get_state(cart_position,tip_speed)
        observation, reward, done, info = env.step(action)
        new_tip_speed = observation[1]
        new_cart_pos = observation[0]
        new_cart_state, new_state = get_state(new_cart_pos, new_tip_speed)

        # Update Q-table
        update_q(state, action, new_state, cart_state, new_cart_state, reward)

        episode_reward += reward

        if done:
            #print("Episode finished after {} timesteps".format(timestep+1))
            break
    if i_episode > 0 & i_episode < 4900:
        epsilon = epsilon - epsilon/i_episode
    rewards.append(episode_reward)
print(Q)
print(epsilon)
env.render()
env.close()

print(rewards)
plt.plot(running_mean(rewards,N))
plt.show()