#!/usr/bin/env python

"""
Wait command for all workers join to a master.  It returns 0, or 1 on
error.  It checks responses from HTTP UI (http://master:8080/).  It
retries in 600 seconds by default.
"""

import sys
import urllib2
import time
import argparse
import logging
import itertools
import xml.etree.ElementTree as ET

FORMAT = '%(filename)s-%(levelname)s: %(message)s'
logging.basicConfig(format=FORMAT)
logger = logging.getLogger(__name__)

class WaitStart():
    """Wait command until all workers connect to a master."""

    __parser = argparse.ArgumentParser()
    __parser.add_argument('workers', type=int,
                          help='number of workers to wait')
    __parser.add_argument('--terse', action='store_true',
                          help=argparse.SUPPRESS)
    __parser.add_argument('--host', type=str, default='localhost')
    __parser.add_argument('--port', type=int, default=8080)
    __parser.add_argument('--interval', type=int, default=10)
    __parser.add_argument('--retrys', type=int, default=60)

    @classmethod
    def __parse_args(cls):
        """Parses arguments."""

        args = cls.__parser.parse_args(sys.argv[1:])
        if args.port < 0:
            cls.__parser.error('port is too small')
        if args.port > 65535:
            cls.__parser.error('port is too large')
        if args.retrys < 0:
            cls.__parser.error('retrys should be positive')
        args.terse = False
        return args

    @classmethod
    def get_alive_workers_field_from_response(cls, html):
        """
        Gets the alive workers field from a string HTML.  It expects
        '
Alive Workers: 1
'.  It returns 0
        when no fields exist.
        """
        elems0 = ET.fromstring(html).findall('.//li')
        elems = filter(lambda x: len(x) > 0 and x[0].text == 'Alive Workers:',
                       elems0)
        if len(elems) == 0:
            logger.warning('No Alive Workers field.')
            return 0
        elif len(elems[0]) == 0:
            raise Error('(never)')
        elif not elems[0][0].tail.strip().isdigit():
            logger.error('Bad Alive Workers field: {}.'
                         .format(elems[0][0].text))
            raise ValueError('Bad Alive Workers field.')
        else:
            return int(elems[0][0].tail)

    @classmethod
    def run(cls):
        """Waits until all workers become available."""

        args = cls.__parse_args()

        if args.terse:
            logger.setLevel(logging.WARNING)
        else:
            logger.setLevel(logging.INFO)

        url = 'http://{}:{}/'.format(args.host, args.port)
        req = urllib2.Request(url)

        logger.info('Wait for workers to join: workers={}.'
                    .format(args.workers))
        if args.retrys == 0:
            repeater = itertools.repeat(True)
        else:
            repeater = itertools.repeat(True, args.retrys)
        for _ in repeater:
            try:
                response = urllib2.urlopen(req)
            except urllib2.URLError as e:
                logger.warning('Bad connection: {}.'.format(str(e)))
                time.sleep(args.interval)
                continue
            except urllib2.HTTPError as e:
                logger.warning('Bad response: {}.'.format(str(e)))
                time.sleep(args.interval)
                continue
            logger.info('Master response: code={}.'.format(response.getcode()))
            body = response.read()
            ##logger.debug('Master response: {}.'.format(body))
            try:
                nworkers = cls.get_alive_workers_field_from_response(body)
            except ValueError as e:
                logger.error('Bad response: {}.'.format(str(e)))
                logger.exception(e)
                return 1

            logger.info('Current workers: workers={}.'.format(nworkers))
            if nworkers >= args.workers:
                logger.info('All expected workers joined to master.')
                return 0
            time.sleep(args.interval)
        logger.error('Retry limit reached.')
        return 1

if __name__ == '__main__':
    v = WaitStart.run()
    sys.exit(v)


from torchvision import transforms
from torch.utils.data import dataset, dataloader
from torchvision.datasets.folder import default_loader
from utils.RandomErasing import RandomErasing
from utils.RandomSampler import RandomSampler
from PIL import Image
from opt import opt
import os
import re
import numpy as np

class Data():
    def __init__(self):
        train_transform = transforms.Compose([
            transforms.Resize((384, 128), interpolation=3),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            RandomErasing(probability=0.5, mean=[0.0, 0.0, 0.0])
        ])

        test_transform = transforms.Compose([
            transforms.Resize((384, 128), interpolation=3),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        self.trainset = Market1501(train_transform, 'train', opt.data_path)
        self.testset = Market1501(test_transform, 'test', opt.data_path)
        self.queryset = Market1501(test_transform, 'query', opt.data_path)

        self.train_loader = dataloader.DataLoader(self.trainset,
                                                  sampler=RandomSampler(self.trainset, batch_id=opt.batchid,
                                                                        batch_image=opt.batchimage),
                                                  batch_size=opt.batchid * opt.batchimage, num_workers=8,
                                                  pin_memory=True)#8

        self.test_loader = dataloader.DataLoader(self.testset, batch_size=opt.batchtest, num_workers=8, pin_memory=True)#8
        self.query_loader = dataloader.DataLoader(self.queryset, batch_size=opt.batchtest, num_workers=8,
                                                  pin_memory=True)#8

        if opt.mode == 'vis':
            self.query_image = test_transform(default_loader(opt.query_image))

class Market1501(dataset.Dataset):
    def __init__(self, transform, dtype, data_path):

        self.transform = transform
        self.loader = default_loader
        self.data_path = data_path

        if dtype == 'train':
            self.data_path += '/bounding_box_train'
        elif dtype == 'test':
            self.data_path += '/bounding_box_test'
        else:
            self.data_path += '/query'

        self.imgs = [path for path in self.list_pictures(self.data_path) if int(path.split('/')[-1].split('_')[0])  != -1]

        num_pids, num_imgs, num_cams, = self.get_imagedata_info(self.imgs)
        print("Dataset statistics:")
        print("  ----------------------------------------")
        print("  subset   | # ids | # images | # cameras")
        print("  ----------------------------------------")
        print("  ", dtype ," |", str(num_pids) , "| ", str(num_imgs) ," | ", str(num_cams) )
        print("  ----------------------------------------")

        self.ids, self.cameras = self.get_id_cam()
        self.unique_ids = sorted(set(self.ids))
        self._id2label = {_id: idx for idx, _id in enumerate(self.unique_ids)} # Build a dictionary where the key is the id and the element is the label
        # print(self._id2label)

    def __getitem__(self, index):
        # please write the __getitem__ method
        img = Image.open(self.imgs[index-1]) #3*64*128
        target = self._id2label[self.get_id(self.imgs[index-1])]
        if self.transform is not None:
            img = self.transform(img)
        return img, target

    def __len__(self):
        # please write the __len__ method
        return len(self.imgs)

    def get_id_cam(self):
        """"
        :return: person id list and camera list corresponding to dataset image paths
        """
        return [int(file_path.split('/')[-1].split('_')[0]) for file_path in self.imgs], [int(file_path.split('/')[-1].split('_')[1][1]) for file_path in self.imgs]

    def get_id(self,file_path):
        """
        :param file_path: unix style file path
        :return: person id
        """
        return int(file_path.split('/')[-1].split('_')[0])

    def list_pictures(self, directory, ext='jpg|jpeg|bmp|png|ppm|npy'):
        assert os.path.isdir(directory), 'dataset does not exists!{}'.format(directory)
        return sorted([os.path.join(root, f).replace('\\','/')
                       for root, _, files in os.walk(directory) for f in files if re.match(r'([\w]+\.(?:' + ext + '))', f)])

    def get_imagedata_info(self, data):
        pids, cams = [], []

        for file_path in data:
            camid = int(file_path.split('/')[-1].split('_')[1][1])
            pid=int(file_path.split('/')[-1].split('_')[0])
            pids += [pid]
            cams += [camid]
        pids = set(pids)
        cams = set(cams)
        num_pids = len(pids)
        num_cams = len(cams)
        num_imgs = len(data)
        return num_pids, num_imgs, num_cams


import argparse
import inspect
import itertools
from abc import ABC, abstractmethod
from typing import Callable, List

import numpy as np
import onnx
from onnx import helper, numpy_helper
from onnx.onnx_ml_pb2 import GraphProto, ModelProto, NodeProto
from tensorboard.compat.proto.attr_value_pb2 import AttrValue
from tensorboard.compat.proto.graph_pb2 import GraphDef
from tensorboard.compat.proto.node_def_pb2 import NodeDef
from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
from tensorboard.compat.proto.versions_pb2 import VersionDef
from torch.utils.tensorboard import SummaryWriter

_DTYPE_ONNX_TO_TF = {
    1: 1,  # FLOAT
    2: 4,  # UINT8
    3: 6,  # INT8
    4: 17,  # UINT16
    5: 5,  # INT16
    6: 3,  # INT32
    7: 9,  # INT64
    8: 7,  # STRING
    9: 10,  # BOOL
    10: 19,  # FLOAT16
    11: 2,  # DOUBLE
    12: 22,  # UINT32
    13: 23,  # UINT64
    14: 8,  # COMPLEX64
    15: 18,  # COMPLEX128
    16: 14,  # BFLOAT16
}

def dtype_onnx_to_tf(onnx_dtype: int) -> int:
    return _DTYPE_ONNX_TO_TF.get(onnx_dtype, 0)

def get_node_by_input(graph: GraphProto, input_name: str) -> NodeProto:
    """Get the first node in the graph node list that consumes the input_name as the first input.
    Return None if no such node is found.
    """

    for node in graph.node:
        if len(node.input) >= 1 and node.input[0] == input_name:
            return node
    return None

def get_prefix(name: str) -> str:
    return name[: name.rfind("/") + 1]

def parse(graph: GraphProto) -> GraphDef:
    nodes = []

    def _add_io_node(node, type):
        shape_proto = TensorShapeProto(
            dim=[TensorShapeProto.Dim(size=d.dim_value, name=d.dim_param) for d in node.type.tensor_type.shape.dim]
        )
        nodes.append(
            NodeDef(
                name=node.name.encode(encoding="utf_8"),
                op=type,
                input=[],
                attr={
                    "dtype": AttrValue(type=dtype_onnx_to_tf(node.type.tensor_type.elem_type)),
                    "shape": AttrValue(shape=shape_proto),
                },
            )
        )

    for node in graph.input:
        _add_io_node(node, "Input")
    for node in graph.output:
        _add_io_node(node, "Output")

    for node in graph.initializer:
        shape_proto = TensorShapeProto(dim=[TensorShapeProto.Dim(size=d) for d in node.dims])
        data = np.array2string(numpy_helper.to_array(node), separator=",").replace("\n", "").replace(" ", "")
        # Truncate the data if it is too long.
        data = data[:128] + "..." if len(data) > 128 else data
        nodes.append(
            NodeDef(
                name=node.name.encode(encoding="utf_8"),
                op="Const",
                input=[],
                attr={
                    "dtype": AttrValue(type=dtype_onnx_to_tf(node.data_type)),
                    "shape": AttrValue(shape=shape_proto),
                    "data": AttrValue(s=data.encode(encoding="utf_8")),
                },
            )
        )

    value_info_map = {}
    for node in graph.value_info:
        shape_proto = TensorShapeProto(
            dim=[TensorShapeProto.Dim(size=d.dim_value, name=d.dim_param) for d in node.type.tensor_type.shape.dim]
        )
        value_info_map[node.name] = (shape_proto, node.type.tensor_type.elem_type)

    for node in graph.node:
        _attr = []
        for s in node.attribute:
            _attr.append(" = ".join([str(f[1]) for f in s.ListFields()]))
        attr = ", ".join(_attr).encode(encoding="utf_8")
        shape_proto = None
        elem_type = 0
        if node.output[0] in value_info_map:
            shape_proto, elem_type = value_info_map[node.output[0]]
        nodes.append(
            NodeDef(
                name=node.output[0].encode(encoding="utf_8"),
                op=node.op_type,
                input=node.input,
                attr={
                    "dtype": AttrValue(type=dtype_onnx_to_tf(elem_type)),
                    "shape": AttrValue(shape=shape_proto),
                    "parameters": AttrValue(s=attr),
                },
            )
        )

    return GraphDef(node=nodes, versions=VersionDef(producer=22))

def add_onnx_model(self: SummaryWriter, model: ModelProto) -> None:
    self._get_file_writer().add_onnx_graph(parse(model.graph))

# Monkey patch SummaryWriter to add onnx graph.
SummaryWriter.add_onnx_model = add_onnx_model

class TransformerBase(ABC):
    """Abstract base class for transformers that need to be applied to the ONNX graph before
    converting it to TensorBoard. This base class is also used as a registry for all non-abstract transformers.
    NOTE: Transformers are applied in the order they are registered. Adding a new transformer needs to consider
    the dependency between it and existing transformers.
    """

    _TRANSFORMERS = []  # noqa: RUF012

    @classmethod
    def register_transformer(cls, klass):
        if not inspect.isabstract(klass):
            cls._TRANSFORMERS.append(klass())

    @classmethod
    def run_transformers(cls, graph):
        for transformer in cls._TRANSFORMERS:
            transformer.transform(graph)

    def __init_subclass__(cls):
        super().__init_subclass__()
        TransformerBase.register_transformer(cls)

    @abstractmethod
    def transform(self, graph: GraphProto) -> None:
        pass

    def _apply_to_all_names(self, graph: GraphDef, func: Callable[[str], str]) -> None:
        for node in graph.node:
            for idx, name in enumerate(node.input):
                node.input[idx] = func(name)
            for idx, name in enumerate(node.output):
                node.output[idx] = func(name)
        for node in itertools.chain(graph.input, graph.output, graph.initializer, graph.value_info):
            node.name = func(node.name)

class HierarchicalNameTransformer(TransformerBase):
    """TensorBoard can show the graph in different levels for better visualization. PyTorch exporter can already
    add hierarchical information to the ONNX graph nodes and edges, but the behavior is not consistent. For example,
    for graph inputs and outputs, the exported name is _original_module.A.B.C, but for activations and initializers,
    the exported name is /_original_module/A/B/C. This transformer will make the exported names consistent
    to use "/" as the separator. The top level _original_module will also be removed.
    """

    def __init__(self):
        super().__init__()
        self.sections = set()
        self.original_module_name = "_original_module"

    def _add_sections(self, name: str) -> None:
        if "/" in name:
            for sec in name.split("/"):
                if len(sec) > 0:
                    self.sections.add(sec)

    def _get_sections(self, curr_name: str, sections: List[str]) -> None:
        for section in self.sections:
            if curr_name.startswith(section) and (len(curr_name) == len(section) or curr_name[len(section)] == "."):
                sections.append(section)
                if curr_name == section:
                    return
                self._get_sections(curr_name[len(section) + 1 :], sections)
                return
        sections.append(curr_name)

    def _transform_name(self, name: str) -> str:
        if "/" in name:
            if name.startswith(f"/{self.original_module_name}/"):
                name = name[len(self.original_module_name) + 2 :]
            if name.startswith("/"):
                name = name[1:]
            return name

        sections = []
        self._get_sections(name, sections)
        if sections[0] == self.original_module_name:
            sections = sections[1:]
        return "/".join(sections)

    def transform(self, graph: GraphProto) -> None:
        for node in graph.node:
            for name in itertools.chain(node.input, node.output):
                self._add_sections(name)
        for initializer in graph.initializer:
            self._add_sections(initializer.name)
        self._apply_to_all_names(graph, self._transform_name)

class ReplaceNameTransformer(TransformerBase):
    """Abstract class for transformers that need to replace the names of nodes and edges in the ONNX graph.
    The sub-class need to prepare the mapping between old names and new names in the _generate_new_names method.
    """

    def __init__(self):
        super().__init__()
        self.new_names = {}

    def _transform_name(self, name: str) -> str:
        return self.new_names.get(name, name)

    @abstractmethod
    def _generate_new_names(self, graph: GraphProto) -> None:
        pass

    def transform(self, graph: GraphProto) -> None:
        self._generate_new_names(graph)
        self._apply_to_all_names(graph, self._transform_name)

class SplitSqueezeTransformer(ReplaceNameTransformer):
    """Handle the edge names without hierarchical information that are introduced by GatherToSplitFusion."""

    def _generate_new_names(self, graph: GraphProto) -> None:
        for node in graph.node:
            if node.doc_string == "Split for Fused Gather nodes":
                squeeze_node = get_node_by_input(graph, node.output[0])
                assert squeeze_node is not None
                prefix = get_prefix(squeeze_node.output[0])
                if len(prefix) > 0:
                    for output in node.output:
                        self.new_names[output] = prefix + output
                        self.new_names[output + "_grad"] = prefix + output + "_grad"

class ExtraOutputsTransformer(ReplaceNameTransformer):
    """Handle the output names without hierarchical information that are introduced by some fusions.
    This kind of fusions add extra outputs to the node, we will try to extract the hierarchical information
    from the first output of the node and add it to those extra outputs.
    """

    def __init__(self):
        super().__init__()
        self.supported_ops = [
            "LayerNormalization",
            "ConcatTraining",
            "BitmaskDropout",
            "BiasSoftmaxDropout",
            "BitmaskBiasDropout",
        ]

    def _generate_new_names(self, graph: GraphProto) -> None:
        for node in graph.node:
            if node.op_type in self.supported_ops and len(node.output) > 1:
                prefix = get_prefix(node.output[0])
                if len(prefix) > 0:
                    for output in node.output[1:]:
                        if "/" not in output:
                            self.new_names[output] = prefix + output

class YieldOpTransformer(ReplaceNameTransformer):
    """Handle the output names without hierarchical information that are introduced by YieldOp."""

    def _generate_new_names(self, graph: GraphProto) -> None:
        for node in graph.node:
            if node.op_type == "YieldOp":
                # TODO: the length of input and output can be not equal if attribute non_differentiable_outputs
                # is not empty. Need to handle this case.
                assert len(node.input) == len(node.output)
                for idx, input in enumerate(node.input):
                    prefix = get_prefix(input)
                    if len(prefix) > 0:
                        self.new_names[node.output[idx]] = prefix + node.output[idx]

class ListUnpackTransformer(TransformerBase):
    """Tensorboard's NodeDef doesn't have output field, it assumes each node has only one output and put this output as
    node name. If an ONNX node has more than one outputs, we will add a special node typed ListUnpack for each output.
    """

    def __init__(self):
        super().__init__()
        self.ops = {}

    def transform(self, graph: GraphProto) -> None:
        new_nodes = []
        for node in graph.node:
            if len([output for output in node.output if len(output) > 0]) > 1:
                idx = self.ops.get(node.op_type, 0)
                self.ops[node.op_type] = idx + 1
                new_output = f"{get_prefix(node.output[0])}{node.op_type}_{idx!s}_output"
                for output in node.output:
                    if len(output) > 0:
                        new_nodes.append(helper.make_node("ListUnpack", [new_output], [output]))
                node.ClearField("output")
                node.output.extend([new_output])
        if len(new_nodes) > 0:
            graph.node.extend(new_nodes)

class AppendOpTypeTransformer(ReplaceNameTransformer):
    """Tensorboard's node search tool can only index the node name, not the node type. To make the op type searchable,
    add the op type to the beginning of the node name. From the tensorboard, we can also easily get the op type
    even part of the node name is truncated.
    """

    def _generate_new_names(self, graph: GraphProto) -> None:
        for node in graph.node:
            pos = node.output[0].rfind("/")
            self.new_names[node.output[0]] = f"{node.output[0][: pos + 1]}{node.op_type}::{node.output[0][pos + 1 :]}"

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--logdir", help="Tensorboard log directory")
    parser.add_argument("--model", help="ONNX model path")

    args = parser.parse_args()
    if not args.logdir or not args.model:
        print("Fail to convert. Please specify the tensorboard log directory and the onnx model path.")
        parser.print_help()
        return

    model = onnx.load(args.model)
    TransformerBase.run_transformers(model.graph)
    with SummaryWriter(args.logdir) as writer:
        writer.add_onnx_model(model)

    print("Successfully converted the onnx model to tensorboard log directory. Start tensorboard to browse.")

if __name__ == "__main__":
    main()


import os
import cv2
import torch
import random
import pathlib
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams.update({'figure.max_open_warning': 0})

def mask_step(x, step):
    """
    Create a mask to only contain the step-th element starting from the first element. Used to downsample
    """
    mask = np.zeros_like(x)
    mask[::step] = 1
    return mask.astype(bool)

def downsample(df, step):
    """
    Downsample data by the given step. Example, SDD is recorded in 30 fps, with step=30, the fps of the resulting
    df will become 1 fps. With step=12 the result will be 2.5 fps. It will do so individually for each unique
    pedestrian (metaId)
    :param df: pandas DataFrame - necessary to have column 'metaId'
    :param step: int - step size, similar to slicing-step param as in array[start:end:step]
    :return: pd.df - downsampled
    """
    mask = df.groupby(['metaId'])['metaId'].transform(mask_step, step=step)
    return df[mask]

def filter_short_trajectories(df, threshold):
    """
    Filter trajectories that are shorter in timesteps than the threshold
    :param df: pandas df with columns=['x', 'y', 'frame', 'trackId', 'sceneId', 'metaId']
    :param threshold: int - number of timesteps as threshold, only trajectories over threshold are kept
    :return: pd.df with trajectory length over threshold
    """
    len_per_id = df.groupby(by='metaId', as_index=False).count(
    )  # sequence-length for each unique pedestrian
    idx_over_thres = len_per_id[len_per_id['frame'] >= threshold]  # rows which are above threshold
    idx_over_thres = idx_over_thres['metaId'].unique() # only get metaIdx with sequence-length longer than threshold
    df = df[df['metaId'].isin(idx_over_thres)] # filter df to only contain long trajectories
    return df

def groupby_sliding_window(x, window_size, stride):
    x_len = len(x)
    n_chunk = (x_len - window_size) // stride + 1
    idx = []
    metaId = []
    for i in range(n_chunk):
        idx += list(range(i * stride, i * stride + window_size))
        metaId += ['{}_{}'.format(x.metaId.unique()[0], i)] * window_size
    df = x.iloc()[idx]
    df['newMetaId'] = metaId
    return df

def sliding_window(df, window_size, stride):
    """
    Assumes downsampled df, chunks trajectories into chunks of length window_size. When stride < window_size then
    chunked trajectories are overlapping
    :param df: df
    :param window_size: sequence-length of one trajectory, mostly obs_len + pred_len
    :param stride: timesteps to move from one trajectory to the next one
    :return: df with chunked trajectories
    """
    gb = df.groupby(['metaId'], as_index=False)
    df = gb.apply(groupby_sliding_window, window_size=window_size, stride=stride)
    df['metaId'] = pd.factorize(df['newMetaId'], sort=False)[0]
    df = df.drop(columns='newMetaId')
    df = df.reset_index(drop=True)
    return df

def split_at_fragment_lambda(x, frag_idx, gb_frag):
    """ Used only for split_fragmented() """
    metaId = x.metaId.iloc()[0]
    counter = 0
    if metaId in frag_idx:
        split_idx = gb_frag.groups[metaId]
        for split_id in split_idx:
            x.loc[split_id:, 'newMetaId'] = '{}_{}'.format(metaId, counter)
            counter += 1
    return x

def split_fragmented(df):
    """
    Split trajectories when fragmented (defined as frame_{t+1} - frame_{t} > 1)
    Formally, this is done by changing the metaId at the fragmented frame and below
    :param df: DataFrame containing trajectories
    :return: df: DataFrame containing trajectories without fragments
    """

    gb = df.groupby('metaId', as_index=False)
    # calculate frame_{t+1} - frame_{t} and fill NaN which occurs for the first frame of each track
    df['frame_diff'] = gb['frame'].diff().fillna(value=1.0).to_numpy()
    fragmented = df[df['frame_diff'] != 1.0] # df containing all the first frames of fragmentation
    gb_frag = fragmented.groupby('metaId')  # helper for gb.apply
    frag_idx = fragmented.metaId.unique()  # helper for gb.apply
    df['newMetaId'] = df['metaId']  # temporary new metaId

    df = gb.apply(split_at_fragment_lambda, frag_idx, gb_frag)
    df['metaId'] = pd.factorize(df['newMetaId'], sort=False)[0]
    df = df.drop(columns='newMetaId')
    return df

def rot(df, image, k=1):
    '''
    Rotates image and coordinates counter-clockwise by k * 90° within image origin
    :param df: Pandas DataFrame with at least columns 'x' and 'y'
    :param image: PIL Image
    :param k: Number of times to rotate by 90°
    :return: Rotated Dataframe and image
    '''
    xy = df.copy()
    if image.ndim == 3:
        y0, x0, channels = image.shape
    else:
        y0, x0 = image.shape

    xy.loc()[:, 'x'] = xy['x'] - x0 / 2
    xy.loc()[:, 'y'] = xy['y'] - y0 / 2
    c, s = np.cos(-k * np.pi / 2), np.sin(-k * np.pi / 2)
    R = np.array([[c, s], [-s, c]])
    xy.loc()[:, ['x', 'y']] = np.dot(xy[['x', 'y']], R)
    for i in range(k):
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)

    if image.ndim == 3:
        y0, x0, channels = image.shape
    else:
        y0, x0 = image.shape

    xy.loc()[:, 'x'] = xy['x'] + x0 / 2
    xy.loc()[:, 'y'] = xy['y'] + y0 / 2
    return xy, image

def fliplr(df, image):
    '''
    Flip image and coordinates horizontally
    :param df: Pandas DataFrame with at least columns 'x' and 'y'
    :param image: PIL Image
    :return: Flipped Dataframe and image
    '''
    xy = df.copy()
    if image.ndim == 3:
        y0, x0, channels = image.shape
    else:
        y0, x0 = image.shape

    xy.loc()[:, 'x'] = xy['x'] - x0 / 2
    xy.loc()[:, 'y'] = xy['y'] - y0 / 2
    R = np.array([[-1, 0], [0, 1]])
    xy.loc()[:, ['x', 'y']] = np.dot(xy[['x', 'y']], R)
    image = cv2.flip(image, 1)

    if image.ndim == 3:
        y0, x0, channels = image.shape
    else:
        y0, x0 = image.shape

    xy.loc()[:, 'x'] = xy['x'] + x0 / 2
    xy.loc()[:, 'y'] = xy['y'] + y0 / 2
    return xy, image

def augment_data(data, image_path='data/SDD/train', images={}, image_file='reference.jpg', seg_mask=False, use_raw_data=False):
    '''
    Perform data augmentation
    :param data: Pandas df, needs x,y,metaId,sceneId columns
    :param image_path: example - 'data/SDD/val'
    :param images: dict with key being sceneId, value being PIL image
    :param image_file: str, image file name
    :param seg_mask: whether it's a segmentation mask or an image file
    :return:
    '''
    ks = [1, 2, 3]
    for scene in data.sceneId.unique():
        if use_raw_data:
            scene_name, scene_idx = scene.split("_")
            im_path = os.path.join(
                image_path, scene_name, f"video{scene_idx}", image_file)
        else:
            im_path = os.path.join(image_path, scene, image_file)
        if seg_mask:
            im = cv2.imread(im_path, 0)
        else:
            im = cv2.imread(im_path)
        images[scene] = im
    # data without rotation, used so rotated data can be appended to original df
    data_ = data.copy()
    k2rot = {1: '_rot90', 2: '_rot180', 3: '_rot270'}
    for k in ks:
        metaId_max = data['metaId'].max()
        for scene in data_.sceneId.unique():
            if use_raw_data:
                im_path = os.path.join(
                    image_path, scene_name, f"video{scene_idx}", image_file)
            else:
                im_path = os.path.join(image_path, scene, image_file)
            if seg_mask:
                im = cv2.imread(im_path, 0)
            else:
                im = cv2.imread(im_path)

            data_rot, im = rot(data_[data_.sceneId == scene], im, k)
            # image
            rot_angle = k2rot[k]
            images[scene + rot_angle] = im

            data_rot['sceneId'] = scene + rot_angle
            data_rot['metaId'] = data_rot['metaId'] + metaId_max + 1
            data = pd.concat([data, data_rot], axis=0)

    metaId_max = data['metaId'].max()
    for scene in data.sceneId.unique():
        im = images[scene]
        data_flip, im_flip = fliplr(data[data.sceneId == scene], im)
        data_flip['sceneId'] = data_flip['sceneId'] + '_fliplr'
        data_flip['metaId'] = data_flip['metaId'] + metaId_max + 1
        data = pd.concat([data, data_flip], axis=0)
        images[scene + '_fliplr'] = im_flip

    return data, images

def resize_and_pad_image(images, size, pad=2019):
    """ Resize image to desired size and pad image to make it square shaped and ratio of images still is the same, as
    images all have different sizes.
    """
    for key, im in images.items():
        H, W, C = im.shape
        im = cv2.copyMakeBorder(
            im, 0, pad - H, 0, pad - W, cv2.BORDER_CONSTANT)
        im = cv2.resize(im, (size, size), interpolation=cv2.INTER_AREA)
        images[key] = im

def create_images_dict(unique_scene, image_path, image_file='reference.jpg', use_raw_data=False):
    images = {}
    for scene in unique_scene:
        if image_file == 'oracle.png':
            im = cv2.imread(os.path.join(image_path, scene, image_file), 0)
        else:
            if use_raw_data:
                scene_name, scene_idx = scene.split("_")
                im_path = os.path.join(
                    image_path, scene_name, f"video{scene_idx}", image_file)
            else:
                im_path = os.path.join(image_path, scene, image_file)
            im = cv2.imread(im_path)
        images[scene] = im
    # images channel: blue, green, red
    return images

def load_images(scenes, image_path, image_file='reference.jpg'):
    images = {}
    if type(scenes) is list:
        scenes = set(scenes)
    for scene in scenes:
        if image_file == 'oracle.png':
            im = cv2.imread(os.path.join(image_path, scene, image_file), 0)
        else:
            im = cv2.imread(os.path.join(image_path, scene, image_file))
        images[scene] = im
    return images

def get_varf_table(df, varf_list, obs_len):
    if obs_len:
        print(f'Computing variation fatcor by obs_len')
    else:
        print(f'Computing variation fatcor by obs_len + pred_len')

    df_varfs = df.groupby(['metaId', 'label', 'sceneId']).size().reset_index()[['metaId', 'label', 'sceneId']]
    df_varfs['scene'] = df_varfs.sceneId.apply(lambda x: x.split('_')[0])
    for varf in varf_list:
        df_stats = aggregate_per_varf_value(df, varf, obs_len)
        df_varfs = df_varfs.merge(df_stats[['metaId', varf]], on='metaId')
    return df_varfs

def aggregate_per_varf_value(df, varf, obs_len):
    out = df.groupby('metaId').apply(aggregate_per_varf_value_per_metaId, varf, obs_len)
    df_stats = pd.DataFrame(
        [[idx, item[0], item[1]] for idx, item in out.items()],
        columns=['metaId', varf, 'label'])
    return df_stats

def aggregate_per_varf_value_per_metaId(df_meta, varf, obs_len):
    x = df_meta["x"].values
    y = df_meta["y"].values

    # sanity check
    unique_labels = np.unique(df_meta["label"].values)
    assert len(unique_labels) == 1
    label = unique_labels[0]

    unique_frame_step = (
        df_meta['frame'].shift(periods=-1) - df_meta['frame']).iloc[:-1].unique()
    assert len(unique_frame_step) == 1
    frame_step = unique_frame_step[0]

    op, attr = varf.split('_')

    # take the observed trajectory, or obs + pred
    if not obs_len:
        obs_len = len(x)

    # compute stats
    if attr == 'vel':
        stats_seqs = np.sqrt((x[: obs_len-1] - x[1: obs_len]) ** 2 + \
                             (y[: obs_len-1] - y[1: obs_len]) ** 2) / frame_step
    elif attr == 'acc':
        vel = np.sqrt((x[:obs_len-1] - x[1: obs_len]) ** 2 + \
                      (y[:obs_len-1] - y[1: obs_len]) ** 2) / frame_step
        stats_seqs = (vel[:obs_len-2] - vel[1: obs_len-1]) / frame_step
    elif attr == 'dist':
        stats_seqs = df_meta[:obs_len].dist.apply(
            lambda x: x.min() if not isinstance(x, float) else np.inf)
    elif 'den' in attr:
        stats_seqs = df_meta[:obs_len].dist.apply(
            lambda x: x[x < int(attr[3:])].shape[0] if not isinstance(x, float) else 0)
    else:
        raise ValueError(f'Cannot compute {attr} statistic')

    # take statistic for one sequence
    if op == 'max':
        stats = np.max(stats_seqs)
    elif op == 'avg':
        stats = np.mean(stats_seqs)
    elif op == 'min':
        stats = np.min(stats_seqs)
    elif op == 'abs+max':
        stats = np.max(np.abs(stats_seqs))
    elif op == 'abs+avg':
        stats = np.mean(np.abs(stats_seqs))
    elif op == 'abs+min':
        stats = np.mean(np.abs(stats_seqs))
    elif op == 'tot':
        stats = np.sum(stats_seqs)
    else:
        raise ValueError(f'Cannot compute {op} operation')

    return stats, label

def add_range_column(df, varf, varf_ranges, obs_len, inclusive='both'):
    df_stats = aggregate_per_varf_value(df, varf, obs_len)
    for r in varf_ranges:
        df_stats.loc[df_stats[varf].between(r[0], r[1], inclusive=inclusive), f'{varf}_range'] = f'{r[0]}_{r[1]}'
    df = df.merge(df_stats[['metaId', f'{varf}_range']], on='metaId')
    return df

def convert_df_to_dict(df_gb):
    varf_group_dict = dict()
    for g_range in df_gb.groups.keys():
        df_g = df_gb.get_group(g_range)[['metaId', 'sceneId', 'label']].drop_duplicates()
        assert df_g.metaId.nunique() == df_g.shape[0]
        varf_group_dict[g_range] = df_g.to_dict('list')
    return varf_group_dict

def create_dataset_by_agent_type(
    df, labels, out_dir, statistic_only,
    same_group_size=False, selected_scenes=None):

    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    df_label = df[df.label.isin(labels)]
    df_gb = df_label.groupby(by='label', dropna=True)
    n_total = df_label[df_label.metaId == df_label.metaId.unique()[0]].shape[0]
    print('Statistics:\n', df_gb.count()['metaId'] / n_total)
    print('# total:', (df_gb.count()['metaId'] / n_total).sum())
    if not statistic_only:
        agent_group_dict = convert_df_to_dict(df_gb)
        for agent, agent_group in agent_group_dict.items():
            if same_group_size:
                min_n = min([len(g["metaId"]) for g in agent_group_dict.values()])
                meta_id_mask = reduce_group_size(agent_group, agent, min_n)
                df_varf = df_label[df_label.metaId.isin(agent_group['metaId'][meta_id_mask])]
            else:
                df_varf = df_label[df_label.metaId.isin(agent_group['metaId'])]
            if selected_scenes is None:
                out_path = os.path.join(out_dir, f"{agent}.pkl")
                df_varf.to_pickle(out_path)
            else:
                df_scenes = pd.DataFrame()
                for scene_id in selected_scenes:
                    out_dir_scene = os.path.join(out_dir, scene_id)
                    out_path = os.path.join(out_dir_scene, f'{agent}.pkl')
                    pathlib.Path(out_dir_scene).mkdir(parents=True, exist_ok=True)

                    df_scene = df_varf[df_varf.sceneId == scene_id]
                    df_scenes = pd.concat([df_scenes, df_scene], axis=0)
                    print(f'scene_id = {scene_id}, label = {agent}, #= {df_scene.metaId.unique().shape[0]}')
                    df_scene.to_pickle(out_path)
                out_dir_scene = os.path.join(out_dir, '__'.join(selected_scenes))
                pathlib.Path(out_dir_scene).mkdir(parents=True, exist_ok=True)
                print(f'scene_id = {selected_scenes}, label = {agent}, #= {df_scenes.metaId.unique().shape[0]}')
                df_scenes.to_pickle(os.path.join(out_dir_scene, f'{agent}.pkl'))

def create_dataset_given_range(df, varf, varf_ranges, labels, out_dir,
        obs_len, statistic_only, inclusive='both', same_group_size=False):
    """_summary_

    Args:
        df (pd.DataFrame): _description_
        varf (str | list of str): _description_
        varf_ranges (list of tuple | list of list of tuple): _description_
        labels (_type_): _description_
        out_dir (_type_): _description_
        obs_len (_type_): _description_
        statistic_only (bool):
            Whether store the generated dataset or give statistic of categorized dataset only.
        inclusive (str):
            Choices = [both, right, left, neither]
        same_group_size (bool, optional):
            _description_. Defaults to False.
    """
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    df_label = df[df.label.isin(labels)]

    # categorize by factor of variation
    if isinstance(varf_ranges[0], tuple):
        varf = varf[0]
        df_label = add_range_column(df_label, varf, varf_ranges, obs_len, inclusive=inclusive)
        varf_col_name = f'{varf}_range'
    elif isinstance(varf_ranges[0], list):
        for f, r in zip(varf, varf_ranges):
            df_label = add_range_column(df_label, f, r, obs_len, inclusive=inclusive)
        varf_col_name = '__'.join(varf)+'_range'
        nonan_mask = df_label.isna().any(axis=1)
        df_label.loc[~nonan_mask, varf_col_name] = \
            df_label.loc[~nonan_mask, [f + '_range' for f in varf]].agg('__'.join, axis=1)
    else:
        raise ValueError(f'Cannot process {varf}.')
    df_gb = df_label.groupby(by=varf_col_name, dropna=True)
    print('Statistics:\n', df_gb.count()['metaId'].unique().shape[0])
    print('# total:', (df_gb.count()['metaId'].unique().shape[0]).sum())

    # save categorized dataset
    if not statistic_only:
        varf_group_dict = convert_df_to_dict(df_gb)
        for varf_range, varf_group in varf_group_dict.items():
            if same_group_size:
                min_n = min([len(g["metaId"]) for g in varf_group_dict.values()])
                meta_id_mask = reduce_group_size(varf_group, varf_range, min_n)
                df_varf = df_label[df_label.metaId.isin(varf_group['metaId'][meta_id_mask])]
            else:
                df_varf = df_label[df_label.metaId.isin(varf_group['metaId'])]
            out_path = os.path.join(out_dir, f"{varf_range}.pkl")
            df_varf.to_pickle(out_path)

def reduce_group_size(varf_group, varf_range, min_n):
    print(f"Group {varf_range}")
    scene_ids, scene_counts = np.unique(
        varf_group["sceneId"], return_counts=True)
    sorted_unique_scene_counts = np.unique(np.sort(scene_counts))
    total_count = 0
    prev_count = 0
    mask = np.zeros_like(scene_counts).astype(bool)
    for scene_count in sorted_unique_scene_counts:
        total_count += (scene_counts >= scene_count).sum() * \
            (scene_count - prev_count)
        if total_count >= min_n:
            break
        mask[scene_counts == scene_count] = True
        prev_count = scene_count
    total_counts = np.zeros_like(scene_counts)
    total_counts[mask] = scene_counts[mask]
    total_counts[mask == False] = prev_count
    less = True
    while less:
        for i in np.where(mask == False)[0]:
            total_counts[i] += min(1, min_n - total_counts.sum())
            if min_n == total_counts.sum():
                less = False
                break
    varf_group["sceneId"] = np.array(varf_group["sceneId"])
    varf_group["metaId"] = np.array(varf_group["metaId"])
    varf_group["label"] = np.array(varf_group["label"])
    meta_id_mask = np.zeros_like(varf_group["metaId"]).astype(bool)
    for scene_idx, scene_id in enumerate(scene_ids):
        scene_count = total_counts[scene_idx]
        scene_mask = varf_group["sceneId"] == scene_id
        scene_labels = varf_group["label"][scene_mask]
        unique_scene_labels, scene_labels_count = np.unique(
            scene_labels, return_counts=True)
        scene_labels_chosen = []
        while len(scene_labels_chosen) < scene_count:
            for label_idx, (unique_scene_label, scene_label_count) in enumerate(zip(unique_scene_labels, scene_labels_count)):
                if scene_label_count != 0:
                    scene_labels_chosen.append(unique_scene_label)
                    scene_labels_count[label_idx] -= 1
                    if len(scene_labels_chosen) == scene_count:
                        break
        labels_chosen, labels_chosen_count = np.unique(
            scene_labels_chosen, return_counts=True)
        for label, label_count in zip(labels_chosen, labels_chosen_count):
            meta_id_idx = np.where(np.logical_and(
                varf_group["label"] == label, varf_group["sceneId"] == scene_id))[0][:label_count]
            meta_id_mask[meta_id_idx] = True
    return meta_id_mask

def compute_distance_with_neighbors(df_scene):
    return df_scene.apply(lambda_distance_with_neighbors, axis=1, df_scene=df_scene)

def lambda_distance_with_neighbors(row, df_scene, step=12):
    # start = datetime.datetime.now()
    frame_diff = df_scene.frame - row.frame
    df_sim = df_scene[(frame_diff < step/2) & \
        (frame_diff >= -step/2) & (df_scene.metaId != row.metaId)]
    dist = np.inf if df_sim.shape[0] == 0 else compute_distance_xy(df_sim, row.x, row.y)
    # duration = datetime.datetime.now() - start
    # print(f'### meta_id = {row.metaId}, time = {duration}')
    return dist

def compute_distance_xy(df_sim, x, y):
    x_diff = df_sim['x'] - x
    y_diff = df_sim['y'] - y
    dist = np.sqrt((x_diff ** 2 + y_diff ** 2))
    return np.array(dist)

def plot_varf_histograms(df_varf, out_dir):
    stats_all = np.array([])
    varf = df_varf.columns[-1]
    # Visualize data
    for label, indices in df_varf.groupby('label').groups.items():
        if label not in ["Pedestrian", "Biker"]:
            continue
        stats = df_varf.iloc[indices].loc[:, varf].values
        plot_histogram(stats, f'{label}_{varf}', out_dir)
        stats_all = np.append(stats_all, stats)
    plot_histogram(stats_all, f"Mixed_{varf}", out_dir)

def plot_varf_hist_obs_and_complete(df_varf, out_dir):
    varf_obs, varf_com = df_varf.columns[-2], df_varf.columns[-1]
    data_all_diff, data_all_obs, data_all_com = np.array([]), np.array([]), np.array([])
    # Visualize data
    for label, indices in df_varf.groupby('label').groups.items():
        if label not in ["Pedestrian", "Biker"]:
            continue
        data_obs = df_varf.iloc[indices].loc[:, varf_obs].values
        data_com = df_varf.iloc[indices].loc[:, varf_com].values
        data_diff = data_obs - data_com
        plot_histogram(data_diff, f'{label}_{varf_obs}_element_diff', out_dir)
        plot_histogram_overlay(data_obs, data_com, f'{label}_{varf_obs}_distr_diff', out_dir)
        data_all_diff = np.append(data_all_diff, data_diff)
        data_all_obs = np.append(data_all_obs, data_obs)
        data_all_com = np.append(data_all_com, data_com)
    plot_histogram(data_all_diff, f"Mixed_{varf_obs}_element_diff", out_dir)
    plot_histogram_overlay(data_all_obs, data_all_com, f'Mixed_{varf_obs}_distr_diff', out_dir)

def plot_histogram(data, title, out_dir, format='png'):
    fig = plt.figure()
    data, stats = filter_long_tail_arr(data)
    mean, std, min, max, p_zero, p_filter = stats
    sns.histplot(data, kde=True)
    plt.title(
        f"{title}, \nMean: {mean}, Std: {std}, Min: {min}, Max: {max}, Zero: {p_zero}, Filter: {p_filter}")
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    plt.savefig(os.path.join(out_dir, title+'.'+format))
    plt.close(fig)

def plot_histogram_overlay(data_obs, data_com, title, out_dir, format='png'):
    fig = plt.figure()
    data_obs, _ = filter_long_tail_arr(data_obs)
    data_com, _ = filter_long_tail_arr(data_com)
    data_obs = data_obs[data_obs != 0]
    data_com = data_com[data_com != 0]
    df_obs = pd.DataFrame(data_obs, columns=['value'])
    df_obs['type'] = 'observation'
    df_com = pd.DataFrame(data_com, columns=['value'])
    df_com['type'] = 'complete'
    df_cat = pd.concat([df_obs, df_com], axis=0)
    df_cat = df_cat.reset_index(drop=True)
    sns.histplot(data=df_cat, x='value', hue="type")
    plt.title(title)
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    plt.savefig(os.path.join(out_dir, title+'.'+format))
    plt.close(fig)

def plot_pairplot(df_varfs, varf_list, label, title, out_dir, kind='kde', format='png'):
    if label == 'Mixed':
        df_label = df_varfs[
            (df_varfs.label == 'Pedestrian') | (df_varfs.label == 'Biker')]
    elif label == 'All':
        df_label = df_varfs
    else:
        df_label = df_varfs[df_varfs.label == label]

    fig = plt.figure()
    df_label_filter, p_filter = filter_long_tail_df(
        df_label[['metaId', 'scene', 'label']+varf_list], varf_list)
    plot_kws = dict(s=1) if kind == 'scatter' else None
    sns.pairplot(
        data=df_label_filter,
        hue="scene",
        vars=varf_list,
        plot_kws=plot_kws,
        diag_kind="hist",
        kind=kind
    )
    title = f'{title}_{label}_{kind}_{str(p_filter)}'
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    plt.savefig(os.path.join(out_dir, title+'.'+format))
    plt.close(fig)

def plot_jointplot(df_varfs, varf_list, label, title, out_dir, hue, kind='kde', format='png'):
    if label == 'Mixed':
        df_label = df_varfs[
            (df_varfs.label == 'Pedestrian') | (df_varfs.label == 'Biker')]
    elif label == 'All':
        df_label = df_varfs
    else:
        df_label = df_varfs[df_varfs.label == label]

    for i, varf1 in enumerate(varf_list):
        for j, varf2 in enumerate(varf_list):
            if i < j:
                fig = plt.figure()
                df_label_filter, p_filter = filter_long_tail_df(
                    df_label[['metaId', 'scene', 'label', varf1, varf2]], [varf1, varf2])
                try:
                    sns.jointplot(data=df_label_filter, x=varf1, y=varf2,
                        kind=kind, hue=hue)
                except np.linalg.LinAlgError:
                    kind = 'scatter'
                    sns.jointplot(data=df_label_filter, x=varf1, y=varf2,
                        kind=kind, hue=hue)
                except:
                    print('Error!')
                title_save = f'{title}_{hue}_{label}_{varf1}_{varf2}_{kind}_{str(p_filter)}.{format}'
                pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
                plt.savefig(os.path.join(out_dir, title_save))
                plt.close(fig)

def plot_scene_w_numeric(df_varfs, varf, title, out_dir, format='png'):
    df_filter, p_filter = filter_long_tail_df(
        df_varfs[['metaId', 'scene', 'label', varf]], [varf])
    # filter out scene "quad"
    df_filter = df_filter[df_filter.scene != 'quad']

    unique_scenes = df_filter.scene.unique()
    n_scene = unique_scenes.shape[0]
    fig, axs = plt.subplots(4, n_scene+1,
        figsize=(4*(n_scene+1), 16), sharex=True, sharey=True)
    binwidth = df_filter[varf].max() / 30
    for c, scene in enumerate(unique_scenes):
        data = df_filter[(df_filter.scene == scene)]
        axs[0, c].set_title(unique_scenes[c])
        # pedestrain
        sns.histplot(data=data[data.label == 'Pedestrian'], x=varf,
            ax=axs[0, c], stat='probability', binwidth=binwidth)
        # biker
        sns.histplot(data=data[data.label == 'Biker'], x=varf,
            ax=axs[1, c], stat='probability', binwidth=binwidth)
        # mixed
        sns.histplot(data=data[(data.label == 'Pedestrian') | (data.label == 'Biker')], x=varf,
            ax=axs[2, c], hue='label', stat='probability',
            hue_order=['Biker', 'Pedestrian'], binwidth=binwidth)
        # all
        sns.histplot(data=data, x=varf,
            ax=axs[3, c], stat='probability', binwidth=binwidth)
    axs[0, -1].set_title('All scenes')
    # pedestrain
    sns.histplot(data=df_filter[df_filter.label == 'Pedestrian'], x=varf,
        ax=axs[0, -1], stat='probability', binwidth=binwidth)
    # biker
    sns.histplot(data=df_filter[df_filter.label == 'Biker'], x=varf,
        ax=axs[1, -1], stat='probability', binwidth=binwidth)
    # mixed
    sns.histplot(data=df_filter[(df_filter.label == 'Pedestrian') | (df_filter.label == 'Biker')], x=varf,
        ax=axs[2, -1], hue='label', stat='probability',
        hue_order=['Biker', 'Pedestrian'], binwidth=binwidth)
    # all
    sns.histplot(data=df_filter, x=varf,
        ax=axs[3, -1], stat='probability', binwidth=binwidth)
    axs[0, 0].set_ylabel('Pedestrian')
    axs[1, 0].set_ylabel('Biker')
    axs[2, 0].set_ylabel('Pedestrian + Biker')
    axs[3, 0].set_ylabel('All agent types')
    plt.tight_layout()
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    plt.savefig(os.path.join(out_dir, f'{title}_scene_w_{varf}_{p_filter}_noquad.{format}'))
    plt.close(fig)

def filter_long_tail_arr(arr, n=3):
    # for statistics computing
    n_data = arr.shape[0]
    arr = arr[~np.isnan(arr) & (arr != np.inf)]
    if arr.shape[0]:
        mean = np.round(np.mean(arr), 2)
        std = np.round(np.std(arr), 2)
        min = np.round(np.min(arr), 2)
        max = np.round(np.max(arr), 2)
    else:
        raise ValueError('stats array is empty')
    p_zero = np.round((arr == 0).sum() / n_data, 2)
    arr = arr[
        (arr < mean + n * std) & (arr > mean - n * std) & (arr != 0)]
    p_filter = np.round((n_data - arr.shape[0]) / n_data, 2)
    return arr, (mean, std, min, max, p_zero, p_filter)

def filter_long_tail_series(series, n=3):
    full_index = series.index
    series = series[~series.isnull() & (series != np.inf)]
    if series.shape[0]:
        mean = np.round(series.mean(), 2)
        std = np.round(series.std(), 2)
    else:
        raise ValueError('Series is empty')
    series = series[
        (series < mean + n * std) & (series > mean - n * std) & (series != 0)]
    return full_index.drop(series.index)

def filter_long_tail_df(df_varfs, varf_list, n=3):
    idx_out = pd.Index([])
    for varf in varf_list:
        idx_out = idx_out.append(filter_long_tail_series(df_varfs[varf]))
    idx_out_unique = idx_out.unique()
    df_varfs_filter = df_varfs.drop(idx_out_unique)
    p_filter = round(len(idx_out_unique) / df_varfs.shape[0], 2)
    return df_varfs_filter, p_filter

def split_train_val_test_sequentially(
    data_path, train_files, val_split, test_splits=None,
    shuffle=False, share_val_test=False):
    print(f"Split {train_files} given val_split={val_split}, test_split={test_splits}")
    df_train, df_val, df_test = pd.DataFrame([]), pd.DataFrame([]), pd.DataFrame([])
    for train_file, test_split in zip(train_files, test_splits):
        df_train_ = pd.read_pickle(os.path.join(data_path, train_file))
        df_train_, df_val_, df_test_ = dataset_split_by_ratio(
            df_train_, val_split, test_split, shuffle, share_val_test)
        df_train = pd.concat([df_train, df_train_])
        df_val = pd.concat([df_val, df_val_])
        if df_test_ is not None:
            df_test = pd.concat([df_test, df_test_])
    return df_train, df_val, df_test

def dataset_split_by_ratio(df, val_split, test_split=None,
    shuffle=False, share_val_test=False, given_test_meta_ids=None):
    # idx
    unique_meta_ids = np.unique(df["metaId"])
    if shuffle:
        print('Shuffling data')
        np.random.shuffle(unique_meta_ids)
    n_metaId = unique_meta_ids.shape[0]
    n_val = int(val_split) if val_split > 1 else int(val_split * n_metaId)
    # split
    if test_split is not None:
        n_test = int(test_split) if test_split > 1 else int(test_split * n_metaId)
        if share_val_test:
            print('Share validation and test set')
            n_train = n_metaId - n_test
            train_meta_ids, test_meta_ids = np.split(unique_meta_ids, [n_train])
            if n_val != 0:
                interval = n_test // n_val if n_test // n_val > 1 else 3
                val_meta_ids = test_meta_ids[::interval]
                df_val = reduce_df_meta_ids(df, val_meta_ids)
            else:
                df_val = None
            df_test = reduce_df_meta_ids(df, test_meta_ids)
        else:
            print('Validation and test sets are independent')
            n_train = n_metaId - n_val - n_test
            train_meta_ids, val_meta_ids, test_meta_ids = np.split(unique_meta_ids, [n_train, n_train + n_val])
            if given_test_meta_ids is not None:
                test_meta_ids = given_test_meta_ids
                print('Replaced test set by given test meta ids')
            df_test = reduce_df_meta_ids(df, test_meta_ids)
            df_val = reduce_df_meta_ids(df, val_meta_ids)
    else:
        n_train = n_metaId - n_val
        val_meta_ids, train_meta_ids = np.split(
            unique_meta_ids, [n_train])
        df_test = None
        df_val = reduce_df_meta_ids(df, val_meta_ids)
    df_train = reduce_df_meta_ids(df, train_meta_ids)
    return df_train, df_val, df_test

def reduce_df_meta_ids(df, meta_ids):
    return df[(df["metaId"].values == meta_ids[:, None]).sum(axis=0).astype(bool)]

def dataset_split_given_scenes(data_path, files, scenes):
    print(f"Split {files} given scenes={scenes}")
    df = pd.concat([pd.read_pickle(os.path.join(data_path, file)) for file in files])
    df_selected = df[df.sceneId.isin(scenes)]
    return df_selected

def split_train_val_test_randomly(data_dir, data_filename, val_split, test_split, seed=1):
    """
    Generate train / val / test set randomly.
    It will output train.pkl / val.pkl / test.pkl under the same directory of input data file.
    """
    data_folder = data_filename.replace('.pkl', '')
    out_dir = f'{data_dir}/{data_folder}'
    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)

    df = pd.read_pickle(f'{data_dir}/{data_filename}')
    unique_meta_ids = np.unique(df["metaId"])

    n_data = unique_meta_ids.shape[0]
    n_val = int(val_split) if val_split > 1 else int(val_split * n_data)
    n_test = int(test_split) if test_split > 1 else int(test_split * n_data)
    n_train = n_data - n_val - n_test

    set_random_seeds(seed)
    np.random.shuffle(unique_meta_ids)

    train_meta_ids, val_meta_ids, test_meta_ids = \
        np.split(unique_meta_ids, [n_train, n_train + n_val])
    df_train = reduce_df_meta_ids(df, train_meta_ids)
    df_val = reduce_df_meta_ids(df, val_meta_ids)
    df_test = reduce_df_meta_ids(df, test_meta_ids)
    print(f'# data = {unique_meta_ids.shape[0]}')
    print(f'# train = {train_meta_ids.shape[0]}')
    print(f'# val = {val_meta_ids.shape[0]}')
    print(f'# test = {test_meta_ids.shape[0]}')

    df_train.to_pickle(f'{out_dir}/train.pkl')
    df_val.to_pickle(f'{out_dir}/val.pkl')
    df_test.to_pickle(f'{out_dir}/test.pkl')
    print('Split train/val/test set')

def load_predefined_train_val_test(data_path, batch_size, n_train_batch=None, shuffle=False):
    df_train = pd.read_pickle(f'{data_path}/train.pkl')
    df_val = pd.read_pickle(f'{data_path}/val.pkl')
    df_test = pd.read_pickle(f'{data_path}/test.pkl')
    if n_train_batch is not None:
        n_sample = int(batch_size * n_train_batch)
        unique_train_ids = df_train.metaId.unique()
        n_train = unique_train_ids.shape[0]
        assert n_sample <= n_train, \
            f'Training set size ({n_train}) < Sample size ({n_sample})'
        if shuffle:
            np.random.shuffle(unique_train_ids)
        df_train = reduce_df_meta_ids(df_train, unique_train_ids[:n_sample])
    return df_train, df_val, df_test

def prepare_dataeset(
    data_path, load_data, batch_size, n_train_batch,
    train_files, val_files, val_split, test_splits,
    shuffle, share_val_test, mode='train', show_details=False):
    if load_data == 'predefined':
        print('Loading predefined train/val/test sets')
        df_train, df_val, df_test = load_predefined_train_val_test(data_path,
            batch_size=batch_size, n_train_batch=n_train_batch, shuffle=shuffle)
    else:
        print('Splitting train/val/test sets sequentially')
        if mode == 'train':
            assert train_files is not None, 'No train file is provided'
            assert val_files is not None, 'No val file is provided'
            assert val_split is not None, 'No val split is provided'
            if train_files == val_files:
                df_train, df_val, df_test = split_train_val_test_sequentially(data_path,
                    train_files, val_split, test_splits, shuffle, share_val_test)
            else:
                raise NotImplementedError
            df_train = limit_samples(df_train, n_train_batch, batch_size)
        elif mode == 'eval':
            assert val_files is not None, 'No val file is provided'
            df_train, df_val, df_test = split_train_val_test_sequentially(data_path,
                val_files, val_split, test_splits, shuffle, share_val_test)
        else:
            raise NotImplementedError

    if show_details:
        print(f'train_meta_ids: {df_train.metaId.unique()}')
        print(f'val_meta_ids: {df_val.metaId.unique()}')
        print(f'test_meta_ids: {df_test.metaId.unique()}')

    if mode == 'train':
        if df_train is not None: print(f"df_train: {df_train.shape}; #={df_train.metaId.unique().shape[0]}")
        if df_val is not None: print(f"df_val: {df_val.shape}; #={df_val.metaId.unique().shape[0]}")
    if df_test is not None: print(f"df_test: {df_test.shape}; #={df_test.metaId.unique().shape[0]}")

    return df_train, df_val, df_test

def get_meta_ids_focus(df=None, given_meta_ids=None, given_csv=None, random_n=None):
    if given_meta_ids is not None:
        if isinstance(given_meta_ids, int):
            meta_ids_focus = [given_meta_ids]
        elif isinstance(given_meta_ids, list):
            meta_ids_focus = given_meta_ids
        else:
            raise ValueError(f'Invalid given_meta_ids={given_meta_ids}')
    elif given_csv['path'] is not None:
        path = given_csv['path']
        col1, col2, op = given_csv['name'].split('__')
        n_limited = given_csv['n_limited']
        df_result = pd.read_csv(path)
        if op == 'diff':
            df_result.loc[:, 'diff'] = df_result[col1].values - df_result[col2].values
        elif op == 'abs_diff':
            df_result.loc[:, 'diff'] = np.abs(df_result[col1].values - df_result[col2].values)
        else:
            raise ValueError(f'Invalid op={op}')
        meta_ids_focus = df_result.sort_values(
            by='diff', ascending=False).head(n_limited).metaId.values
    elif random_n is not None:
        unique_meta_ids = df.metaId.unique()
        np.random.shuffle(unique_meta_ids)
        meta_ids_focus = unique_meta_ids[:random_n]
    else:
        meta_ids_focus = df.metaId.unique()
    print('Focusing on meta_ids=', meta_ids_focus)
    return meta_ids_focus

def set_random_seeds(random_seed=0):
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed(random_seed)
    np.random.seed(random_seed)
    random.seed(random_seed)
    cv2.setRNGSeed(random_seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def limit_samples(df, num, batch_size, random_ids=True):
    if num is None:
        return df
    num_total = num * batch_size
    meta_ids = np.unique(df["metaId"])
    if random_ids:
        np.random.shuffle(meta_ids)
    meta_ids = meta_ids[:num_total]
    df = reduce_df_meta_ids(df, meta_ids)
    return df


import os
import cv2
import numpy as np
import random
import matplotlib.pyplot as plt
import lpips

import torch
import torch.nn as nn
import torchvision

from torch.distributions.normal import Normal
from skimage.metrics import structural_similarity as ssim
from model import resnet34
from utils import foveat_img

class Actor(nn.Module):

    def __init__(self):
        super().__init__()
        self.backbone = resnet34(pretrained=True)
        self.fixation_conv = nn.Sequential(nn.Conv2d(64, 16, 3, padding=1),
                                           nn.LeakyReLU(inplace=True),
                                           nn.Conv2d(16, 4, 3, padding=1),
                                           nn.LeakyReLU(inplace=True),
                                           nn.Conv2d( 4, 1, 1, padding=0))

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.blur_conv = nn.Sequential(nn.Linear(64, 16),
                                       nn.LeakyReLU(inplace=True),
                                       nn.Linear(16, 4),
                                       nn.LeakyReLU(inplace=True),
                                       nn.Linear( 4, 1))
        self.distribution = Normal(torch.Tensor([0.0]), torch.Tensor([1.0]))

    def forward(self, x, no_blur=False):

        feats, _ = self.backbone(x)
        feats = nn.functional.interpolate(feats, (335, 447))
        center = nn.functional.softmax(self.fixation_conv(feats).flatten())
        blur_mean = self.blur_conv(self.avgpool(feats).flatten(start_dim=1)) * (no_blur == False)
        blur_mean = nn.functional.sigmoid(blur_mean) * 3.5
        self.distribution = Normal(blur_mean.flatten(), torch.Tensor([1.0]).cuda())

        return center

    def pi(self, s, a):
        probs_center = self.forward(s)
        return probs_center[a]

    def update_weight(self, state, action, blur_p, reward, optimizer):
        loss = (-1.0) * reward * torch.log(self.pi(state, action)) +\
            (-1.0) * reward * self.distribution.log_prob(blur_p)
        # update policy parameter \theta
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def eval(img_name, actor, epoch, resnet, loss_fn_vgg):

    with torch.no_grad():

        img = cv2.imread(img_name)
        if img.shape != (335, 447, 3):
            img = cv2.resize(img, (447, 335))
        img = img[:, :, ::-1]

        orig_state = torch.Tensor(img.copy()).permute(2, 0, 1).unsqueeze(0) / 255
        tr = torchvision.transforms.Normalize((0.485, 0.456, 0.406),
                                              (0.229, 0.224, 0.225))
        orig_state[0] = tr(orig_state[0])
        orig_state = orig_state.cuda()
        orig_preds = resnet(orig_state)
        log_orig_preds = nn.functional.log_softmax(orig_preds, dim=1)

        actor.eval()

        fixation_probs = actor(orig_state)
        action = torch.argmax(fixation_probs).item()
        blur_p = actor.distribution.mean

        fov_img, num_full_res_pixels = foveat_img(img, [(action%447, action//447)], blur_p.item()+7, 3, 1.5)
        state = torch.Tensor(fov_img.copy()).permute(2, 0, 1).unsqueeze(0) / 255
        state[0] = tr(state[0])
        state = state.cuda()
        fov_preds = resnet(state)
        log_fov_preds = nn.functional.log_softmax(fov_preds, dim=1)

        kl_div = nn.functional.kl_div(log_fov_preds, log_orig_preds, reduction='sum', log_target=True).item()

        ssi = ssim(fov_img, img, multichannel=True)

        perceptual_loss = loss_fn_vgg(state.cpu(), orig_state.cpu()).item()

        # cv2.imshow("frame", fov_img[:, :, ::-1])
        # key = cv2.waitKey(100)
        if not os.path.exists(os.path.join("outputs", str(epoch))):
            os.makedirs(os.path.join("outputs", str(epoch)))

        img_name = os.path.basename(img_name).split(".")[0]
        cv2.imwrite(os.path.join("outputs", str(epoch), "{}_{}_{}_{}.jpg".format(img_name, action%447, action//447, blur_p.item()+7)),
            fov_img[:, :, ::-1])

        return kl_div, ssi, perceptual_loss, num_full_res_pixels

if __name__ == '__main__':

    resnet = torchvision.models.resnet34(pretrained=True).cuda()
    resnet.eval()
    loss_fn_vgg = lpips.LPIPS(net='vgg')

    actor = Actor().cuda()
    actor.load_state_dict(torch.load("ckpt/25.pth"))
    img_names = os.listdir(os.path.abspath(os.path.join(__file__ ,"../../dataset/COCO-2017-test-500/data")))
    img_names = [name for name in img_names if name.endswith(".jpg") or name.endswith(".png")]
    metrics = []
    random.shuffle(img_names)
    for img_name in img_names:
        print(img_name)
        img_path = os.path.join(os.path.join(__file__ ,"../../dataset/COCO-2017-test-500/data"), img_name)
        result = eval(img_path, actor, 0, resnet, loss_fn_vgg)
        metrics.append(result)
    print(np.mean(metrics, axis=0))

# if __name__ == '__main__':
#     actor = Actor().cuda()
#     actor.load_state_dict(torch.load("ckpt/999.pth"))
#     img_names = os.listdir("data")
#     img_names = [name for name in img_names if name.endswith(".jpg") or name.endswith(".png")]
#     metrics = []
#     random.shuffle(img_names)
#     for img_name in img_names:
#         img_path = os.path.join("data", img_name)
#         result = eval(img_path, actor, 0)
#         metrics.append(result)
#     print(np.mean(metrics, axis=0))


# -*- coding: utf-8 -*-

# A - Something on It

"""
ラーメン 700 円
トッピング 1 つ 100 円

文字列 S の各文字が各トッピング乗せるかどうか (o で乗せる)
"""
RamenPrice = 700

class ToppingOrNot:
    yes = 'o'
    no = 'x'

def main():
    s = input().strip()
    topping_count = s.count(ToppingOrNot.yes)

    price = RamenPrice + 100 * topping_count
    print(price)

if __name__ == '__main__':
    main()


import itertools
import networkx as nx
from tqdm import tqdm
import logging
import collections

import constants as keys
import services
from supermatch.docs_to_sku.promoted import promoted_match
from supermatch.docs_to_sku.set_match import set_match
from supermatch.docs_to_sku.same_names import exact_name_match

class SKUGraphCreator(services.GenericGraph):
    """ Create a graph with items as vertices and barcodes as edges """

    def __init__(self, id_doc_pairs):
        super().__init__()
        self.sku_graph = nx.Graph()
        self.connected_ids = set()
        self.id_doc_pairs = id_doc_pairs
        self.stages = dict()

    def init_sku_graph(self):
        # add all items as nodes
        logging.info("init sku graph..")
        for doc_id in self.id_doc_pairs.keys():
            self.sku_graph.add_node(doc_id)

    def barcode_match(self):
        barcode_id_pairs = collections.defaultdict(set)
        for doc_id, doc in self.id_doc_pairs.items():
            barcodes = doc.get(keys.BARCODES, [])
            if not barcodes:
                continue

            for code in barcodes:
                barcode_id_pairs[code].add(doc_id)

        logging.info(f"adding_edges using {len(barcode_id_pairs)} barcodes")
        for ids in tqdm(barcode_id_pairs.values()):
            edges = itertools.combinations(ids, 2)
            self.sku_graph.add_edges_from(edges)
            self.connected_ids.update(ids)

        self.stages = {**dict.fromkeys(self.connected_ids, "barcode")}

    def create_graph(self):
        self.init_sku_graph()

        logging.info("barcode match..")
        self.barcode_match()

        logging.info("exact_name_match..")
        exact_name_match(self)

        logging.info("set match..")
        set_match(self)

        logging.info("promoted match..")
        promoted_match(self)

        return self.sku_graph, self.stages

    def single_doc_generator(self):
        unmatched_ids = set(self.id_doc_pairs.keys()).difference(self.connected_ids)
        for doc_id in unmatched_ids:
            if "clone" not in doc_id:
                doc = self.id_doc_pairs.get(doc_id, {})
                yield doc_id, doc

    def get_connected_groups(self) -> list:
        id_groups = self.create_connected_component_groups(self.sku_graph)

        # filter out single item groups
        id_groups = [
            id_group
            for id_group in id_groups
            if all(id in self.connected_ids for id in id_group)
        ]
        return id_groups


# @author: Gautam Patel
# Problem Description URL: https://www.hackerrank.com/ProjectEuler+/project-euler-224-almost-right-angled-triangles-ii/problem

def solve(n):
    cnt = 0
    a = 2
    while True:
        b = (a * a / 2)
        c = ((a * a) + 2) / 2
        if (a + b + c) > n:
            break
        if (a*a + b*b) == (c*c - 1):
            cnt += 1
        a += 2
    return cnt

q = int(input())
for i in range(q):
    n = int(input())
    print(solve(n))


f = open("input.txt", "r")
lines = f.readlines()
lines = [line.replace('\n','') for line in lines]

# part 1

line = lines[0]

for x in range(4,len(line)):
    marker = line[x-4:x]
    if len(set(marker))==4:
        break

print(x)

# part 2

for x in range(14,len(line)):
    marker = line[x-14:x]
    if len(set(marker))==14:
        break

print(x)


#@JejakCheat
""" Userbot module for getting info about any user on Telegram(including you!). """

import os

from telethon.tl.functions.photos import GetUserPhotosRequest
from telethon.tl.functions.users import GetFullUserRequest
from telethon.tl.types import MessageEntityMentionName
from telethon.utils import get_input_location
from userbot import CMD_HELP, TEMP_DOWNLOAD_DIRECTORY
from userbot.events import register

@register(pattern=".siapa(?: |$)(.*)", outgoing=True)
async def who(event):

    await event.edit(
        "`Sedang mencari informasi dari` **Anak Gang Sebelah...**")

    if not os.path.isdir(TEMP_DOWNLOAD_DIRECTORY):
        os.makedirs(TEMP_DOWNLOAD_DIRECTORY)

    replied_user = await get_user(event)

    try:
        photo, caption = await fetch_info(replied_user, event)
    except AttributeError:
        return event.edit("`Tidak dapat mengambil info dari pengguna itu, Karena di privasi .`")

    message_id_to_reply = event.message.reply_to_msg_id

    if not message_id_to_reply:
        message_id_to_reply = None

    try:
        await event.client.send_file(event.chat_id,
                                     photo,
                                     caption=caption,
                                     link_preview=False,
                                     force_document=False,
                                     reply_to=message_id_to_reply,
                                     parse_mode="html")

        if not photo.startswith("http"):
            os.remove(photo)
        await event.delete()

    except TypeError:
        await event.edit(caption, parse_mode="html")

async def get_user(event):
    """ Get the user from argument or replied message. """
    if event.reply_to_msg_id and not event.pattern_match.group(1):
        previous_message = await event.get_reply_message()
        replied_user = await event.client(
            GetFullUserRequest(previous_message.from_id))
    else:
        user = event.pattern_match.group(1)

        if user.isnumeric():
            user = int(user)

        if not user:
            self_user = await event.client.get_me()
            user = self_user.id

        if event.message.entities is not None:
            probable_user_mention_entity = event.message.entities[0]

            if isinstance(probable_user_mention_entity,
                          MessageEntityMentionName):
                user_id = probable_user_mention_entity.user_id
                replied_user = await event.client(GetFullUserRequest(user_id))
                return replied_user
        try:
            user_object = await event.client.get_entity(user)
            replied_user = await event.client(
                GetFullUserRequest(user_object.id))
        except (TypeError, ValueError) as err:
            return await event.edit(str(err))

    return replied_user

async def fetch_info(replied_user, event):
    """ Get details from the User object. """
    replied_user_profile_photos = await event.client(
        GetUserPhotosRequest(user_id=replied_user.user.id,
                             offset=42,
                             max_id=0,
                             limit=80))
    replied_user_profile_photos_count = "Kamu perlu bantuan dalam mengunggah foto profil?."
    try:
        replied_user_profile_photos_count = replied_user_profile_photos.count
    except AttributeError:
        pass
    user_id = replied_user.user.id
    first_name = replied_user.user.first_name
    last_name = replied_user.user.last_name
    try:
        dc_id, location = get_input_location(replied_user.profile_photo)
    except Exception as e:
        dc_id = "Couldn't fetch DC ID!"
        location = str(e)
    common_chat = replied_user.common_chats_count
    username = replied_user.user.username
    user_bio = replied_user.about
    is_bot = replied_user.user.bot
    restricted = replied_user.user.restricted
    verified = replied_user.user.verified
    photo = await event.client.download_profile_photo(user_id,
                                                      TEMP_DOWNLOAD_DIRECTORY +
                                                      str(user_id) + ".jpg",
                                                      download_big=True)
    first_name = first_name.replace(
        "\u2060", "") if first_name else ("Pengguna ini tidak memiliki Nama Depan")
    last_name = last_name.replace(
        "\u2060", "") if last_name else ("Pengguna ini tidak memiliki Nama Belakang")
    username = "@{}".format(username) if username else (
        "Pengguna ini tidak memiliki Nama Pengguna")
    user_bio = "Pengguna ini tidak memiliki BIO" if not user_bio else user_bio

    caption = "INFORMASI USER:\n\n"
    caption += f"Nama Depan: {first_name}\n"
    caption += f"Nama Belakang: {last_name}\n"
    caption += f"Nama Pengguna: {username}\n"
    caption += f"ID Pusat Data: {dc_id}\n"
    caption += f"Jumlah PP: {replied_user_profile_photos_count}\n"
    caption += f"Dia bot: {is_bot}\n"
    caption += f"Di Dibatasi: {restricted}\n"
    caption += f"Verifikasi telegram: {verified}\n"
    caption += f"ID: {user_id}\n"
    caption += f"Grup yang sama: {common_chat}\n\n"
    caption += f"Link Profile Selamanya Milik Pengguna ini: "
    caption += f"{first_name}\n"
    caption += f"Bio: {user_bio}\n\n"
    caption += f"Bot By : #SenturyBot"

    return photo, caption

CMD_HELP.update({
    "siapa":
    ">`.siapa <username> atau balas pesan seseorang dengan mengetik .siapa`"
    "\nFungsi: Mendapatkan Info Pengguna Tersebut."
})


# Code by Lilie Boizumault, for QI Informatique, France
# 2020

# files needed :
# ardSerial.py written by Rongzhong Li, Petoi

# many thanks to David Amos for his great tutorial on Speech Recognition
# https://realpython.com/python-speech-recognition/#installing-speechrecognition

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

# Controlling Nybble with your voice
# Controler Nybble avec ta voix

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

# FIRST STEP : FIND YOUR MICROPHONE
# Uncomment this part, comment the rest

# PREMIERE ETAPE : SELECTIONNER LE BON MICRO
# Decommenter cette partie et commenter le reste

# import speech_recognition as sr

# r = sr.Recognizer()

## utilisation du micro

# affichage de la liste des micro
#print(sr.Microphone.list_microphone_names()) # look for your micro in the list (get the index)
# chercher l'indice du micro voulu

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

# SECOND STEP : TEST YOUR MICROPHONE
# Uncomment this part, comment the rest

# DEUXIEME ETAPE : TESTER LE MICRO
# Decommenter cette partie et commenter le reste

# import speech_recognition as sr

# r = sr.Recognizer()
# mic = sr.Microphone(device_index=0)

# with mic as source :
    # r.adjust_for_ambient_noise(source)
    # audio = r.listen(source)
# print(r.recognize_google(audio, language='fr-FR'))

#####################################################################
# THIRD STEP : SPEECH RECOGNITION AND NYBBLE
# Uncomment this part, comment the rest

# TROISIEME ETAPE : NYBBLE ET LA RECONNAISSANCE VOCALE
# Decommenter cette partie et commenter le reste

## Script de reconnaissance vocale

#!/usr/bin/python
from ardSerial import * #importer le fichier pour communiquer avec Nybble

import speech_recognition as sr #importer la bibliothèque SpeechRecognition

# Création des objets Microphone et Recognizer
recognizer = sr.Recognizer()
microphone = sr.Microphone(device_index=0)

# function to translate the "speech" order to "written" order
def recognize_speech(recognizer, microphone) : # fonction qui traduit l'ordre vocal
    # en ordre écrit
    with microphone as source : # listen to the speech
        recognizer.adjust_for_ambient_noise(source)
        audio = recognizer.listen(source)

    try:
        order = recognizer.recognize_google(audio, language='fr-FR')
    except sr.RequestError: # API unreachable
        # si l'API ne répond pas ou n'est pas joignable
        order = "error1"
        print("API unavailable")
    except sr.UnknownValueError: # unknown speech
        # le mot n'a pas pu être identifié
        order = "error2"
        print("Unable to recognize speech")

    return order

# this function translate the "written" order in an order for Nybble
# WARNING : FOR ENGLISH USERS --> NEED TO CHANGE THE DETECTED WORDS AND THE LANGUAGE DETECTED IN RECOGNIZE_GOOGLE

def translate_order(order) : #fonction qui traduit l'ordre écrit en ordre pour
# Nybble
    order = order.lower() # on s'assure que le texte est en minuscules
    if order == "marche" :
        order_nybble = 'kwk'
    elif order == "debout" :
        order_nybble = 'kbalance'
    elif order == "assis" :
        order_nybble = 'ksit'
    elif order == "bonjour" :
        order_nybble = 'whi_Nybble'
    elif order =="stop" :
        order_nybble = 'kbalance'
    elif order =="repos" :
        order_nybble = 'd'
    else : # si l'ordre n'est pas un mot que l'on a code --> ne rien faire
        print(order)
        order_nybble = ' ' # if the order is none of these orders : don't do anything
    return order_nybble

while(True) :
    last_task = []
    task = []

    order = recognize_speech(recognizer, microphone)
    order_nybble = translate_order(order)

    if not(order_nybble == ' ') : # si on a reçu un ordre compatible avec Nybble
        task = [order_nybble, 0]   # task to be sended to Nybble

    if not(task == last_task) : # si on n'est pas deja en train de faire cette tache
        # if the order is not the current task
        last_task = task
        wrapper(task)
        print(task)
        if (task == ['d', 0]) : break # si on dit "repos" : fin du script
        # if the order is "rest" --> end of script


# -*- coding: utf-8 -*-

n = int(input())
p = list(map(int, input().split()))
a = [list(map(int, list(input()))) for _ in range(n)]

groups = []

for i in range(n):
    s = set()
    s.add(i)
    for j in range(n):
        if a[i][j] == 1:
            s.add(j)
    if len(s) > 1:
        groups.append(s)

i = 0
while i < len(groups):
    j = i + 1
    while j < len(groups):
        if len(groups[i] & groups[j]) == 0:
            j += 1
            continue
        groups[i] |= groups[j]
        groups.pop(j)
        j = i + 1
    i += 1

groups = [sorted(list(s)) for s in groups]

for g in groups:
    i = len(g)
    while i > 0:
        j = 0
        while j < i - 1:
            if p[g[j]] > p[g[j+1]]:
                p[g[j]], p[g[j+1]] = p[g[j+1]], p[g[j]]
            j += 1
        i -= 1

print(' '.join(map(str, p)))


#****************************************************************
"""
### THE BELOW TEXT IS OUTDATED and needs to be updated

#This replaces a previous version of gcard_helper.py by using the HTMLParser class
#This allows for more robust parsing of the html mother directory to find gcard files
#Even better would be to use BeautifulSoup, which would allow for the code to be condensed as:
#```import requests
#from bs4 import BeautifulSoup
#page = requests.get('http://www.website.com')
#bs = BeautifulSoup(page.content, features='lxml')
#for link in bs.findAll('a'):
#    print(link.get('href'))```
#Unfortunately, this module must be imported and cannot be gaurannted that it will be on whatever
#server this software will live on, so it is safer to instead use HTMLParser which is more common
####
#This file takes in a UserSubmissionID and gcard url from db_batch_entry and passes it through
#a few functions to download the gcards from the specified location and to enter them into the
#appropriate gcard table in the database.
# Some effort should be developed to sanitize the gcard files to prevent
# against sql injection attacks
"""
#***************************************************************
from __future__ import print_function
import argparse, subprocess, os
import utils, fs, html_reader
import glob

def Lund_Downloader(url_dir,lund_urls,lund_dir):
    if len(lund_urls) == 0:
        print("No Lund files found (they must end in '{0}'). Is the online repository correct?".format(fs.lund_identifying_text ))
        exit()
    else:
        for url_ending in lund_urls:
            utils.printer('Lund URL name is: '+url_ending)
            lund_text = html_reader.html_reader(url_dir+'/'+url_ending,'')[0]#This returns a tuple, we need the contents of the tuple
            utils.printer2('HTML from lund link is: {0}'.format(lund_text))
            lund_text_db = lund_text.replace('"',"'") #This isn't strictly needed but SQLite can't read " into data fields, only ' characters
            print("\t Gathered lund file '{0}'".format(url_ending))
            filename = lund_dir+"/"+url_ending
            with open(filename,"a") as file: file.write(lund_text_db)

def Lund_Entry(url_dir, target_dir):
    """ Download or copy lund files and return the name of
    the target directory.

    Inputs:
    -------
    url_dir - A string containing the directory or path to lund file(s).

    Returns:
    --------
    lund_dir - A string containing the name of the downloaded directory.

    A few cases can occur:

    1) One local file, extension will be .txt, .dat, or .lund and the string
       will not contain http.
    2) Several local files, no extension will be given.  The string will not
       contain http.
    3) One web file, extension will be .txt, .dat, or .lund and the string
       will contain http.
    4) Many web files, no extension will be given.  The string will contain http.

    """

    lund_extensions = ['.dat', '.txt', '.lund']
    lund_dir = target_dir

    # A case used to work around not downloading for types 1/3
    if url_dir == "no_download":
        print('Not downloading files due to SCard type.')
        return lund_dir

    if os.path.exists(lund_dir):
        print('Lund directory already exists, not downloading again.')
        return lund_dir

    # Create dir.
    subprocess.call(['mkdir', '-p', lund_dir])

    # Case 3/4
    if 'http' in url_dir:

        # Single web file
        if any([ext in url_dir for ext in lund_extensions]):
            lund_dir_unformatted = url_dir.split("/")
            filename = lund_dir_unformatted[len(lund_dir_unformatted)-1]

            lund_text = html_reader.html_reader(url_dir,'')[0]#This returns a tuple, we need the contents of the tuple
            utils.printer2('HTML from lund link is: {0}'.format(lund_text))
            lund_text_db = lund_text.replace('"',"'") #This isn't strictly needed but SQLite can't read " into data fields, only ' characters
            print("\t Gathered lund file '{0}'".format(url_dir))
            with open(lund_dir+"/"+filename,"a") as file:
                    file.write(lund_text_db)

        # Web directory
        else:
            raw_html, lund_urls = html_reader.html_reader(url_dir, fs.lund_identifying_text)
            lund_dir_unformatted = url_dir.split("//")[1]
            subprocess.call(['mkdir','-p',lund_dir])
            Lund_Downloader(url_dir,lund_urls,lund_dir)

    # Case 1/2
    else:

        # Single local file
        if any([ext in url_dir for ext in lund_extensions]):
            subprocess.call(['cp', url_dir, lund_dir + '/'])

        # Local directory, many files
        else:
            print('Downloading all files in {}'.format(url_dir))

            lund_files = glob.glob(url_dir + '*')
            print(lund_files)

            for lf in lund_files:
                if any([ext in lf for ext in lund_extensions]):
                    subprocess.call(['cp', lf, lund_dir+'/'])

    return lund_dir

def count_files(url_dir):
    """
    We need to know how many files are going
    to be downloaded before we do this job.  This
    is used in the queue system.

    Inputs:
    -------
    - url_dir (str) - Specifies the location of the
    lund files.

    Returns:
    --------
    - nfiles (int) - The number of files to be downloaded.

    """
    lund_extensions = ['.dat', '.txt', '.lund']

    # A case used to work around not downloading for types 1/3
    if url_dir == "no_download":
        print('Not downloading files due to SCard type.')
        return 0

    # Case 3/4
    if 'http' in url_dir:

        # Single web file
        if any([ext in url_dir for ext in lund_extensions]):
            return 1

        # Web directory
        else:
            raw_html, lund_urls = html_reader.html_reader(url_dir, fs.lund_identifying_text)
            return len(lund_urls)

    # Case 1/2
    else:

        # Single local file
        if any([ext in url_dir for ext in lund_extensions]):
            return 1

        # Local directory, many files
        else:
            lund_files = glob.glob(url_dir + '*')
            return len(lund_files)

    # Something weird happened.
    return 0


import requests

def sort_parm(parm: {}):
    """
    将字典中的key排序，转换为地址参数字符串，往往API参数签名都需要使用到
    :param parm:
    :return:
    """
    str_parm = ''
    for p in sorted(parm):
        # 每次排完序的加到串中
        # str类型需要转化为url编码格式
        # if isinstance(parm[p], str):
        #     str_parm = f"{str_parm}{p}={parm[p]}&"  # str_parm + str(p) + "=" + str(quote(parm[p])) + "&"
        # continue
        # str_parm = str_parm + str(p) + "=" + str(parm[p]) + "&"
        str_parm = f"{str_parm}{p}={parm[p]}&"

    return str_parm

def request(url: str, method: str, param=None, headers=None, is_json=True, content_type: str = None, time_out=None,
            cookie=None, prox_url=None):
    """
    一个通用的http请求方法，目前只支持get与post,后期会更新
    :param prox_url:
    :param cookie: cookies 采用字典保存
    :param url: 请求的地址
    :param method: 请求的方法
    :param param:   请求参数，一个字典对象
    :param headers: 请求头，一个字典对象
    :param is_json: 返回的结果是否转换成json对象
    :param content_type: 请求时的Content-type类型，比如application/json
    :param time_out: 超时
    :return: 返回文本或json
    """
    try:
        prox = None
        if prox_url:
            prox = {
                'http': prox_url,
                'https': prox_url,
            }
        if method == 'get':
            r = requests.get(url=url, params=param, headers=headers, timeout=time_out, cookies=cookie, proxies=prox)
            r.encoding = r.apparent_encoding  # 服务器传过来的编码格式需要 先转换一下，再转成json格式 否则json格式的数据存在乱码
            result = r.json() if is_json else r.text  # 相当于问题表达式，意思是如果is_json不真，result=r.json()否则result = r.text
            return result
        elif method == 'post':
            if content_type == 'application/json':
                headers['Content-type'] = "application/json;charset=UTF-8"
                r = requests.post(url=url, json=param, headers=headers, timeout=time_out, cookies=cookie, proxies=prox)
                r.encoding = r.apparent_encoding
                result = r.json() if is_json else r.text
                return result
            else:
                r = requests.post(url=url, data=param, headers=headers, timeout=time_out, cookies=cookie, proxies=prox)
                r.encoding = r.apparent_encoding
                result = r.json() if is_json else r.text
                return result
        else:
            print("http method not allowed")
    except Exception as e:
        print("http请求报错:{0}".format(e))

def getText(url: str, params_obj={}, headers={}) -> str:
    """
    使用get方式请求地址
    :param url: 请求地址
    :param params_obj: 请求头，一个字典对象
    :param headers: 请求参数，一个字典对象
    :return: 返回文本
    """
    return request(url, "get", params_obj, headers, False)

def getJson(url: str, params_obj={}, headers={}) -> str:
    """
    使用get方式请求地址
    :param url: 请求地址
    :param params_obj: 请求参数，一个字典对象
    :param headers:  请求头，一个字典对象
    :return: 返回字典类型JSON对象
    """
    return request(url, "get", params_obj, headers, True)

def postText(url: str, params_obj={}, headers={}):
    """
    使用post方式请求地址
    :param url: 请求地址
    :param params_obj: 请求参数，一个字典对象
    :param headers: 请求头，一个字典对象
    :return: 返回文本
    """
    return request(url, "post", params_obj, headers, False)

def postJson(url: str, params_obj={}, headers={}):
    """
    使用post方式请求地址
    :param url: 请求地址
    :param params_obj: 请求参数，一个字典对象
    :param headers: 请求头，一个字典对象
    :return: 返回字典类型JSON对象
    """
    return request(url, "post", params_obj, headers, True)

def postJsonContent(url: str, content_obj={}, headers={}):
    """
    使用post方式请求地址,请求提交的内容为json
    :param url: 请求地址
    :param content_obj: 请求参数，一个字典对象
    :param headers: 请求头，一个字典对象
    :return: 返回字典类型JSON对象
    """
    return request(url, "post", content_obj, headers, True, "application/json")

def postFormContent(url: str, content_obj={}, headers={}, prox_url=None):
    """
    使用post方式请求地址,请求提交的内容为json
    :param prox_url:
    :param url: 请求地址
    :param content_obj: 请求参数，一个字典对象
    :param headers: 请求头，一个字典对象
    :return: 返回字典类型JSON对象
    """
    return request(url, "post", content_obj, headers, True, "application/x-www-form-urlencoded", prox_url=prox_url)