# -*- coding: utf-8 -*-
# vim: sw=4:ts=4:expandtab
"""
riko.modules.fetchpage
~~~~~~~~~~~~~~~~~~~~~~
Provides functions for fetching web pages.

Fetches the source of a given web site as a string. This data can then be
converted into an RSS feed or merged with other data in your Pipe using the
`regex` module.

Examples:
    basic usage::

        >>> from riko.modules.fetchpage import pipe
        >>> from riko import get_path
        >>> from meza.compat import decode
        >>>
        >>> url = get_path('cnn.html')
        >>> conf = {'url': url, 'start': '<title>', 'end': '</title>'}
        >>> resp = next(pipe(conf=conf))['content'][:21]
        >>> decode(resp) == 'CNN.com International'
        True

Attributes:
    OPTS (dict): The default pipe options
    DEFAULTS (dict): The default parser options
"""
import pygogo as gogo

from meza.compat import encode

from . import processor
from riko.bado import coroutine, return_value, io
from riko.parsers import get_text
from riko.utils import betwix, fetch, get_abspath

OPTS = {"ftype": "none"}
logger = gogo.Gogo(__name__, monolog=True).logger

def get_string(content, start, end):
    # TODO: convert relative links to absolute
    # TODO: remove the closing tag if using an HTML tag stripped of HTML tags
    # TODO: clean html with Tidy
    content = encode(content)
    start_pos = content.find(encode(start)) if start else 0
    right = content[start_pos + (len(start) if start else 0) :]
    end_pos = right[1:].find(encode(end)) + 1 if end else len(right)
    return right[:end_pos] if end_pos > 0 else right

@coroutine
def async_parser(_, objconf, skip=False, **kwargs):
    """Asynchronously parses the pipe content

    Args:
        _ (None): Ignored
        objconf (obj): The pipe configuration (an Objectify instance)
        skip (bool): Don't parse the content
        kwargs (dict): Keyword arguments

    Kwargs:
        assign (str): Attribute to assign parsed content (default: content)
        stream (dict): The original item

    Returns:
        Iter[dict]: The stream of items

    Examples:
        >>> from riko import get_path
        >>> from riko.bado import react
        >>> from riko.bado.mock import FakeReactor
        >>> from meza.fntools import Objectify
        >>> from meza.compat import decode
        >>>
        >>> def run(reactor):
        ...     callback = lambda x: print(decode(next(x)['content'][:32]))
        ...     url = get_path('cnn.html')
        ...     conf = {'url': url, 'start': '<title>', 'end': '</title>'}
        ...     objconf = Objectify(conf)
        ...     kwargs = {'stream': {}, 'assign': 'content'}
        ...     d = async_parser(None, objconf, **kwargs)
        ...     return d.addCallbacks(callback, logger.error)
        >>>
        >>> try:
        ...     react(run, _reactor=FakeReactor())
        ... except SystemExit:
        ...     pass
        ...
        CNN.com International - Breaking
    """
    if skip:
        stream = kwargs["stream"]
    else:
        url = get_abspath(objconf.url)
        content = yield io.async_url_read(url)
        parsed = get_string(content, objconf.start, objconf.end)
        detagged = get_text(parsed) if objconf.detag else parsed
        splits = detagged.split(objconf.token) if objconf.token else [detagged]
        stream = ({kwargs["assign"]: chunk} for chunk in splits)

    return_value(stream)

def parser(_, objconf, skip=False, **kwargs):
    """Parses the pipe content

    Args:
        _ (None): Ignored
        objconf (obj): The pipe configuration (an Objectify instance)
        skip (bool): Don't parse the content

    Returns:
        Iter[dict]: The stream of items

    Examples:
        >>> from meza.fntools import Objectify
        >>> from riko import get_path
        >>> from meza.compat import decode
        >>>
        >>> url = get_path('cnn.html')
        >>> conf = {'url': url, 'start': '<title>', 'end': '</title>'}
        >>> objconf = Objectify(conf)
        >>> kwargs = {'stream': {}, 'assign': 'content'}
        >>> result = parser(None, objconf, **kwargs)
        >>> resp = next(result)['content'][:21]
        >>> decode(resp) == 'CNN.com International'
        True
    """
    if skip:
        stream = kwargs["stream"]
    else:
        with fetch(decode=True, **objconf) as f:
            sliced = betwix(f, objconf.start, objconf.end, True)
            content = "\n".join(sliced)

        parsed = get_string(content, objconf.start, objconf.end)
        detagged = get_text(parsed) if objconf.detag else parsed
        splits = detagged.split(objconf.token) if objconf.token else [detagged]
        stream = ({kwargs["assign"]: chunk} for chunk in splits)

    return stream

@processor(isasync=True, **OPTS)
def async_pipe(*args, **kwargs):
    """A source that asynchronously fetches the content of a given web site as
    a string.

    Args:
        item (dict): The entry to process
        kwargs (dict): The keyword arguments passed to the wrapper

    Kwargs:
        conf (dict): The pipe configuration. Must contain the key 'url'. May
            contain the keys 'start', 'end', 'token', or 'detag'.

            url (str): The web site to fetch
            start (str): The starting string to fetch (exclusive, default:
                None).

            end (str): The ending string to fetch (exclusive, default: None).
            token (str): The tokenizer delimiter string (default: None).
            detag (bool): Remove html tags from content (default: False).

        assign (str): Attribute to assign parsed content (default: content)

    Returns:
        dict: twisted.internet.defer.Deferred item

    Examples:
        >>> from riko import get_path
        >>> from riko.bado import react
        >>> from riko.bado.mock import FakeReactor
        >>> from meza.compat import decode
        >>>
        >>> resp = 'html PUBLIC "-//W3C//DTD XHTML+RDFa 1.0//EN" "'
        >>> def run(reactor):
        ...     callback = lambda x: print(decode(next(x)['content']) == resp)
        ...     url, path = get_path('bbc.html'), 'value.items'
        ...     conf = {'url': url, 'start': 'DOCTYPE ', 'end': 'http'}
        ...     d = async_pipe(conf=conf)
        ...     return d.addCallbacks(callback, logger.error)
        >>>
        >>> try:
        ...     react(run, _reactor=FakeReactor())
        ... except SystemExit:
        ...     pass
        ...
        True
    """
    return async_parser(*args, **kwargs)

@processor(**OPTS)
def pipe(*args, **kwargs):
    """A source that fetches the content of a given web site as a string.

    Args:
        item (dict): The entry to process
        kwargs (dict): The keyword arguments passed to the wrapper

    Kwargs:
        conf (dict): The pipe configuration. Must contain the key 'url'. May
            contain the keys 'start', 'end', 'token', or 'detag'.

            url (str): The web site to fetch
            start (str): The starting string to fetch (exclusive, default:
                None).

            end (str): The ending string to fetch (exclusive, default: None).
            token (str): The tokenizer delimiter string (default: None).
            detag (bool): Remove html tags from content (default: False).

        assign (str): Attribute to assign parsed content (default: content)

    Yields:
        dict: item

    Examples:
        >>> from riko import get_path
        >>> from meza.compat import decode
        >>>
        >>> url = get_path('bbc.html')
        >>> conf = {'url': url, 'start': 'DOCTYPE ', 'end': 'http'}
        >>> content = 'html PUBLIC "-//W3C//DTD XHTML+RDFa 1.0//EN" "'
        >>> decode(next(pipe(conf=conf))['content']) == content
        True
    """
    return parser(*args, **kwargs)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/5/2 9:48
# @Author  : yaomy

# real+render+fuse 所有物体
import json
import torch
import random
import torch.utils.data
import torch.nn.functional as F
import torchvision.transforms as transforms
import pickle as pkl
import numpy as np
import yaml
import os
import pickle
from pathlib import Path
from PIL import Image
import cv2
import numpy.ma as ma
from mmcv import Config
from lib.transform.coordinate import crop_resize_by_warp_affine

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJETCT_PATH = Path(os.path.realpath(__file__)).parent.parent.parent
CONFIG = Config.fromfile(f'{PROJETCT_PATH}/config/linemod/lm_v3.py')

class PoseDataset(torch.utils.data.Dataset):
    def __init__(self, mode, num_point, add_noise, root, noise_trans, num_kps, cls_type=None, cfg=CONFIG):
        self.obj_dict = {
            'ape': 1, 'benchvise': 2, 'bowl': 3, 'cam': 4, 'can': 5, 'cat': 6, 'cup': 7, 'driller': 8,
            'duck': 9, 'eggbox': 10, 'glue': 11, 'holepuncher': 12, 'iron': 13, 'lamp': 14, 'phone': 15,
        }

        self.obj_name = {v: k for k, v in self.obj_dict.items()}

        if cls_type is None or cls_type == 'all':
            self.objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
        else:
            self.objlist = [self.obj_dict[cls_type]]

        print(f'{mode} Use Object: ', self.objlist)
        self.mode = mode
        self.num_point = num_point
        self.num_kps = num_kps
        self.cfg = cfg

        self.back = cfg.Data.BACK
        self.num_pt_mesh = num_point
        self.num_pt_mesh_large = 2600
        self.num_pt_mesh_small = num_point

        self.noise_trans = noise_trans
        self.add_noise = add_noise
        self.height = 480.0
        self.width = 640.0
        self.scale = 1000.0

        self.ymap = np.array([[j for i in range(640)] for j in range(480)])
        self.xmap = np.array([[i for i in range(640)] for j in range(480)])
        self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
        # 要归一化后才能用这个
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        self.sym_obj = [7, 8]
        self.obj_dict_name = {v: k for k, v in self.obj_dict.items()}

        self.intrinsic_matrix = {
            'linemod': np.array([[572.4114, 0., 325.2611],
                                 [0., 573.57043, 242.04899],
                                 [0., 0., 1.]]),
            'blender': np.array([[700., 0., 320.],
                                 [0., 700., 240.],
                                 [0., 0., 1.]]),
        }

        self.ori_iamge = np.zeros((480, 640))
        self.gt_pose = torch.zeros((4, 4))
        self.root = root
        self.backgrounds = []
        self.synthetic_all = []
        self.real_all = []
        self.diameter = []
        self.pt = {}
        self.extent = {}
        self.lf_border = {}
        self.meta = {}
        self.kps = {}
        self.centers = {}
        info_yml = open(f'{CURRENT_DIR}/dataset_config/models_info.yml', 'r')
        info = yaml.load(info_yml, Loader=yaml.FullLoader)

        with open(f'{self.root}/test/test_bboxes/bbox_yolov3_all.json', 'r') as f:
            eval_bbox_all = json.load(f)

        for item in self.objlist:
            if self.mode == 'train':
                cls_type = self.obj_dict_name[item]

                if item == 4:
                    cls_type = 'camera'

                # real_img_pth = os.path.join(cls_root, "train.txt")
                render_root = os.path.join(self.root, 'train', f'{item:06d}')
                real_root = os.path.join(self.root, 'test', f'{item:06d}')
                real_img_pth = os.path.join(self.root, 'image_set', f"{cls_type}_train.txt")
                real_lst = read_lines(real_img_pth)
                rnd_paths = sorted(Path(f'{render_root}/rgb').glob('*.png'))

                with open(f'{render_root}/scene_gt.json', 'r') as f:
                    render_gt = json.load(f)
                with open(f'{real_root}/scene_gt.json', 'r') as f:
                    real_gt= json.load(f)

                with open(f'{render_root}/scene_gt_info.json', 'r') as f:
                    render_gt_info = json.load(f)
                with open(f'{real_root}/scene_gt_info.json', 'r') as f:
                    real_gt_info = json.load(f)

                with open(f'{render_root}/scene_camera.json', 'r') as f:
                    render_gt_camera = json.load(f)
                with open(f'{real_root}/scene_camera.json', 'r') as f:
                    real_gt_camera = json.load(f)

                self.synthetic_all += [{
                    'cls_id': item, 'im_id': syn.stem,
                    'cam_R_m2c': render_gt[f'{int(syn.stem)}'][0]['cam_R_m2c'],
                    'cam_t_m2c': render_gt[f'{int(syn.stem)}'][0]['cam_t_m2c'],
                    'bbox': render_gt_info[f'{int(syn.stem)}'][0]['bbox_visib'],
                    'cam_K': render_gt_camera[f'{int(syn.stem)}']['cam_K'],
                    # 'depth_scale': render_gt_camera[f'{int(syn.stem)}']['depth_scale'],
                    'depth_scale': 1000.,
                    'type': 'syn'
                } for syn in rnd_paths]

                self.real_all += [{
                    'cls_id': item, 'im_id': real_id,
                    'cam_R_m2c': real_gt[f'{int(real_id)}'][0]['cam_R_m2c'],
                    'cam_t_m2c': real_gt[f'{int(real_id)}'][0]['cam_t_m2c'],
                    'bbox': real_gt_info[f'{int(real_id)}'][0]['bbox_visib'],
                    'cam_K': real_gt_camera[f'{int(real_id)}']['cam_K'],
                    # 'depth_scale': real_gt_camera[f'{int(real_id)}']['depth_scale'],
                    'depth_scale': 1000.,
                    'type': 'real'
                } for real_id in real_lst]

                self.backgrounds += sorted(Path(self.back).glob('*.jpg'))
            else:
                tst_img_pth = os.path.join(self.root, 'image_set', f"{cls_type}_test.txt")
                real_lst = read_lines(tst_img_pth)
                real_root = os.path.join(self.root, 'test', f'{item:06d}')
                with open(f'{real_root}/scene_gt.json', 'r') as f:
                    real_gt= json.load(f)
                with open(f'{real_root}/scene_gt_info.json', 'r') as f:
                    real_gt_info = json.load(f)
                with open(f'{real_root}/scene_camera.json', 'r') as f:
                    real_gt_camera = json.load(f)
                self.real_all += [{
                    'cls_id': item, 'im_id': real_id,
                    'cam_R_m2c': real_gt[f'{int(real_id)}'][0]['cam_R_m2c'],
                    'cam_t_m2c': real_gt[f'{int(real_id)}'][0]['cam_t_m2c'],
                    'bbox': real_gt_info[f'{int(real_id)}'][0]['bbox_visib'],
                    'eval_bbox': eval_bbox_all[f'{int(item)}/{int(real_id)}'][0]['bbox_est'],
                    'cam_K': real_gt_camera[f'{int(real_id)}']['cam_K'],
                    # 'depth_scale': real_gt_camera[f'{int(real_id)}'][0]['depth_scale'],
                    'depth_scale': 1000.,
                    'type': 'real'
                } for real_id in real_lst]

            self.pt[item], self.extent[item], self.lf_border[item] = self._load_model('{0}/models/obj_{1}.ply'.format(self.root, '%06d' % item))
            self.diameter.append(info[item]['diameter'] / 1000.0)

            kps_pattern = os.path.join(CURRENT_DIR, 'dataset_config', 'kps_orb9_fps',
                                       f"{self.obj_name[item]}_8_kps.txt")
            corner_pattern = os.path.join(CURRENT_DIR, 'dataset_config', 'kps_orb9_fps',
                                          f"{self.obj_name[item]}_corners.txt")

            kps = np.loadtxt(kps_pattern, dtype=np.float32)
            corners = np.loadtxt(corner_pattern, dtype=np.float32)
            # center = np.mean(corners, axis=0).reshape(1, 3)
            center = np.mean(self.pt[item], axis=0).reshape(1, 3)
            self.kps[item] = kps
            self.centers[item] = center

        self.all_lst = self.real_all + self.synthetic_all

        # self.all_lst = self.synthetic_all

        print(f'len of real {self.mode}: ', len(self.real_all))
        print(f'len of synthetic {self.mode}: ', len(self.synthetic_all))
        print(f'len of background {self.mode}: ', len(self.backgrounds))
        print(f'len of all {self.mode}: ', len(self.all_lst), '\n')

    def _load_real_data(self, item):
        im_id = item['im_id']
        cls_id = item['cls_id']
        if item['type'] == 'syn':
            cls_root = os.path.join(self.root, 'train', f'{cls_id:06d}')
        else:
            cls_root = os.path.join(self.root, 'test', f'{cls_id:06d}')

        with Image.open(os.path.join(cls_root, "depth/{}.png".format('%06d' % int(im_id)))) as di:
            depth = np.array(di)

        with open(f'{cls_root}/normal/{int(im_id):06d}-normal.pkl', 'rb') as f_n:
            normal = pickle.load(f_n)

        with open(f'{cls_root}/xyz/{int(im_id):06d}-coordinate.pkl', 'rb') as f_c:
            coordinate = pickle.load(f_c)

        with open(f'{cls_root}/xyz/{int(im_id):06d}-region.pkl', 'rb') as f_r:
            region = pickle.load(f_r)

        # if self.mode == 'eval':
        #     with Image.open(
        #             os.path.join(self.root, "segnet_results", f'{int(cls_id):02d}_label/{int(im_id):04d}_label.png')
        #     ) as li:
        #         label = np.array(li)
        # else:
        with Image.open(os.path.join(cls_root, "mask_visib/{}_000000.png".format('%06d' % int(im_id)))) as li:
            label = np.array(li)

        with Image.open(os.path.join(cls_root, "rgb/{}.png".format('%06d' % int(im_id)))) as ri:
            # if self.add_noise:
            #     ri = self.trancolor(ri)
            img = np.array(ri)[:, :, :3]

            if item['type'] == 'render':
                with Image.open(os.path.join(cls_root, "mask/{}_000000.png".format('%06d' % int(im_id)))) as li:
                    label_back = np.array(li)
                mask_back = ma.getmaskarray(ma.masked_equal(label_back, 0))
                seed = random.choice(self.backgrounds)
                with Image.open(seed) as bf:
                    background = np.asarray(bf.convert("RGB"))

                back = cv2.resize(np.array(self.trancolor(background)), (self.width, self.height))
                img = back * mask_back + img

        self.ori_image = img.copy()
        cam_scale = item['depth_scale']

        target_r = np.resize(np.array(item['cam_R_m2c']), (3, 3))
        target_t = np.array(item['cam_t_m2c']) / cam_scale
        self.gt_pose = self.rt2matrix(target_r, target_t)
        if self.mode == 'eval':
            bbox = item['eval_bbox']
        else:
            bbox = item['bbox']

        mask_label = np.ma.getmaskarray(np.ma.masked_equal(label, np.array(255)))
        # mask_depth = np.ma.getmaskarray(np.ma.masked_not_equal(depth, 0))
        # mask_label = mask_label * mask_depth

        return {
            'img': img,
            'depth': depth,
            'normal': normal,
            'coordinate': coordinate,
            'region': region,
            # 'croped_normal': croped_normal,
            'mask_label': mask_label,
            'cls_id': cls_id,
            'K': np.asarray(item['cam_K']).reshape(3, 3),
            'bbox': bbox,
            'cam_scale': cam_scale,
            'target_r': target_r,
            'target_t': target_t,
            'type': item['type'],
        }

    def _load_data(self, item):
        while True:
            ds = self._load_real_data(item)
            if ds is None:
                print(f'Real data error pass ! {item["type"]}')
                item = random.choice(self.all_lst)
            else:
                break

        out_size = self.cfg.Data.OUT_SIZE
        input_size = self.cfg.Data.INPUT_SIZE
        resize_type = self.cfg.Data.RESIZE_TYPE
        rmin, rmax, cmin, cmax = get_square_bbox(ds['bbox'])

        img_masked = np.array(ds['img'])[rmin:rmax, cmin:cmax, :]

        cam_fx, cam_fy, cam_cx, cam_cy = ds['K'][0, 0], ds['K'][1, 1], ds['K'][0, 2], ds['K'][1, 2]

        cls_id = ds['cls_id']
        target_r, target_t = ds['target_r'], ds['target_t']

        if ds['normal'].shape == (self.height, self.width, 3):
            normal_ = ds['normal'][rmin:rmax, cmin:cmax, :]
        else:
            normal_ = ds['normal']
        hole_normal = np.ones((int(self.height), int(self.width), 3))
        hole_normal[rmin:rmax, cmin:cmax, :] = normal_
        normal = hole_normal * (hole_normal != [1, 1, 1])

        if resize_type == 'resize':
            resize_normal = cv2.resize(normal, (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        else:
            resize_normal= crop_resize_by_warp_affine(
                normal, [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), out_size, interpolation=cv2.INTER_NEAREST
            )

        # croped_normal_ = ds['croped_normal']
        # croped_normal = croped_normal_ * (croped_normal_ != [1, 1, 1])
        #
        # if resize_type == 'resize':
        #     resize_croped_normal = cv2.resize(croped_normal, (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        # else:
        #     resize_croped_normal= crop_resize_by_warp_affine(
        #         croped_normal, [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), 64, interpolation=cv2.INTER_NEAREST
        #     )

        if ds['coordinate'].shape == (self.height, self.width, 3):
            coordinate = ds['coordinate'][rmin:rmax, cmin:cmax, :]
        else:
            coordinate = ds['coordinate']

        # 变成[0, 1]
        # coordinate_map = np.ones_like(coordinate)
        hole_coordinate = np.zeros((int(self.height), int(self.width), 3))
        hole_coordinate[rmin:rmax, cmin:cmax, :] = coordinate
        coordinate_map = np.zeros_like(hole_coordinate)
        lf_border = self.lf_border[cls_id]
        extent = self.extent[cls_id]

        mask_obj = (hole_coordinate != [0, 0, 0])[..., 0]
        mask = (ds['mask_label']*mask_obj).astype(np.float32)

        if resize_type == 'resize':
            resize_mask = cv2.resize(mask[:, :, None], (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        else:
            resize_mask = crop_resize_by_warp_affine(
                mask[:, :, None], [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), out_size,
                interpolation=cv2.INTER_NEAREST
            )

        choose = mask[rmin:rmax, cmin:cmax, None].flatten().nonzero()[0]
        if len(choose) > self.num_point:
            c_mask = np.zeros(len(choose), dtype=int)
            c_mask[:self.num_point] = 1
            np.random.shuffle(c_mask)
            # num_points+8
            choose = choose[c_mask.nonzero()]
        else:
            try:
                choose = np.pad(choose, (0, self.num_point - len(choose)), 'wrap')
            except ValueError:
                print('choose error:', item)
                return None
        choose = np.array([choose])

        resize_choose = resize_mask.flatten().nonzero()[0]
        if len(resize_choose) > self.num_point:
            c_mask = np.zeros(len(resize_choose), dtype=int)
            c_mask[:self.num_point] = 1
            np.random.shuffle(c_mask)
            # num_points+8
            resize_choose = resize_choose[c_mask.nonzero()]
        else:
            try:
                resize_choose = np.pad(resize_choose, (0, self.num_point - len(resize_choose)), 'wrap')
            except ValueError:
                print('choose error:', item)
                return None

        resize_choose = np.array([resize_choose])

        if resize_type == 'resize':
            hole_coordinate_map = cv2.resize(hole_coordinate, (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        else:
            hole_coordinate_map = crop_resize_by_warp_affine(
                hole_coordinate, [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), out_size,
                interpolation=cv2.INTER_NEAREST
            )

        coordinate_choosed = hole_coordinate_map.reshape(-1, 3)[resize_choose[0]].T

        coordinate_map[:, :, 0] = (hole_coordinate[:, :, 0] - lf_border[0]) / extent[0]
        coordinate_map[:, :, 1] = (hole_coordinate[:, :, 1] - lf_border[1]) / extent[1]
        coordinate_map[:, :, 2] = (hole_coordinate[:, :, 2] - lf_border[2]) / extent[2]
        coordinate_map = coordinate_map * (hole_coordinate != [0, 0, 0])

        if resize_type == 'resize':
            resize_coordinate_map = cv2.resize(coordinate_map, (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        else:
            resize_coordinate_map = crop_resize_by_warp_affine(
                coordinate_map, [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), out_size,
                interpolation=cv2.INTER_NEAREST
            )

        if ds['region'].shape == (self.height, self.width, 3):
            region = ds['region'][rmin:rmax, cmin:cmax]
        else:
            region = ds['region']
        hole_region = np.zeros((int(self.height), int(self.width)))
        hole_region[rmin:rmax, cmin:cmax] = region
        if resize_type == 'resize':
            resize_region = cv2.resize(hole_region, (out_size, out_size), interpolation=cv2.INTER_NEAREST)
        else:
            resize_region = crop_resize_by_warp_affine(
                hole_region, [(cmin + cmax) / 2, (rmin + rmax) / 2], int(rmax - rmin), out_size,
                interpolation=cv2.INTER_NEAREST
            )

        ori_kps = self.kps[cls_id] / 1.0
        # ori_center = self.centers[cls_id] / 1.0
        ori_center = np.array([[0., 0., 0.]])
        all_kps = np.concatenate([ori_kps, ori_center], axis=0)

        kps = np.dot(ori_kps, target_r.T) + target_t
        center = np.dot(ori_center, target_r.T) + target_t
        kps = np.concatenate([kps, center], axis=0)
        u = kps[:, 0] * cam_fx / kps[:, 2] + cam_cx
        v = kps[:, 1] * cam_fy / kps[:, 2] + cam_cy

        # [8, 3]
        uvd1 = np.stack([u, v, kps[:, 2]], axis=1)

        model_points = self.pt[cls_id]
        model_points:np.array
        target = model_points @ target_r.T + target_t
        add_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for _ in range(3)])

        if self.add_noise:
            target = np.add(target, add_t)

        if resize_type == 'resize':
            reize_img = cv2.resize(img_masked/255., (input_size, input_size), interpolation=cv2.INTER_LINEAR)
        else:
            reize_img = crop_resize_by_warp_affine(
                ds['img']/255., [(cmin+cmax)/2, (rmin+rmax)/2], int(rmax-rmin), input_size, interpolation=cv2.INTER_LINEAR
            )

        kps_scale = float(rmax-rmin) / self.cfg.Data.INPUT_SIZE
        resize_scale = float(rmax-rmin) / self.cfg.Data.OUT_SIZE

        resize_uvd = uvd1.copy()
        resize_uvd[..., 0] -= cmin
        resize_uvd[..., 1] -= rmin
        resize_uvd[..., :2] /= kps_scale

        resize_xmap = torch.LongTensor([[i for i in range(out_size)] for _ in range(out_size)]) / out_size
        resize_ymap = torch.LongTensor([[j for _ in range(out_size)] for j in range(out_size)]) / out_size

        resize_cam_fx, resize_cam_fy, resize_cam_cx, resize_cam_cy = cam_fx / resize_scale, cam_fy / resize_scale, out_size / 2, out_size / 2

        if resize_type == 'resize':
            x_map = cv2.resize(self.xmap[rmin:rmax, cmin:cmax], (out_size, out_size), cv2.INTER_LINEAR)/640.
            y_map = cv2.resize(self.ymap[rmin:rmax, cmin:cmax], (out_size, out_size), cv2.INTER_LINEAR)/480.
        else:
            x_map = crop_resize_by_warp_affine(
                self.xmap/640., [(cmin+cmax)/2, (rmin+rmax)/2], int(rmax-rmin), out_size, interpolation=cv2.INTER_NEAREST
            )

            y_map = crop_resize_by_warp_affine(
                self.ymap/480., [(cmin+cmax)/2, (rmin+rmax)/2], int(rmax-rmin), out_size, interpolation=cv2.INTER_NEAREST
            )

        return {
            'target': torch.from_numpy(target.astype(np.float32)),  # [Np, 3]
            'model_points': torch.from_numpy(model_points.astype(np.float32)),
            'cls_id': torch.LongTensor([self.objlist.index(cls_id)]),
            'choose': torch.LongTensor(choose.astype(np.int32)),
            'intrinsic': torch.FloatTensor([cam_fx, cam_fy, cam_cx, cam_cy]),
            'uvd1': torch.from_numpy(uvd1),  # [9, 3]
            'trans_kps': torch.from_numpy(kps.astype(np.float32)),
            'kps': torch.from_numpy(ori_kps.astype(np.float32)),
            'all_kps': torch.from_numpy(all_kps.astype(np.float32)),
            'trans_center': torch.from_numpy(center.astype(np.float32)),
            'center': torch.from_numpy(ori_center.astype(np.float32)),
            'target_r': torch.from_numpy(target_r.astype(np.float32)),
            'target_t': torch.from_numpy(target_t.astype(np.float32)),
            'bbox': torch.FloatTensor([rmin, rmax, cmin, cmax]),
            'type': ds['type'],
            'xmap': torch.from_numpy(x_map.astype(np.float32)).unsqueeze(dim=0),
            'ymap': torch.from_numpy(y_map.astype(np.float32)).unsqueeze(dim=0),
            'ori_mask': torch.from_numpy(ds['mask_label'].astype(np.float32)),
            'ori_depth': torch.from_numpy(ds['depth'].astype(np.float32)),
            'ori_coordinate': torch.from_numpy(hole_coordinate.astype(np.float32)),
            'coordinate_choosed': torch.from_numpy(coordinate_choosed.astype(np.float32)),
            'resize_img': torch.from_numpy(reize_img.astype(np.float32)).permute(2, 0, 1),
            'resize_normal': F.normalize(torch.from_numpy(resize_normal).permute(2, 0, 1), p=2, dim=0).float(),
            'resize_uvd': torch.from_numpy(resize_uvd),
            'resize_scale': torch.tensor(resize_scale),
            'kps_scale': torch.tensor(kps_scale),
            'resize_xmap': resize_xmap,
            'resize_ymap': resize_ymap,
            'resize_choose': torch.LongTensor(resize_choose.astype(np.int32)),
            'resize_intrinsic': torch.FloatTensor([resize_cam_fx, resize_cam_fy, resize_cam_cx, resize_cam_cy]),
            'resize_coordinate': torch.from_numpy((resize_coordinate_map*resize_mask[:, :, None]).astype(np.float32)).permute(2, 0, 1),
            'extent': torch.from_numpy(extent),  # [3]
            'lfborder': torch.from_numpy(lf_border),  # [3]
            # 'resize_croped_normal': F.normalize(torch.from_numpy(resize_croped_normal).permute(2, 0, 1), p=2, dim=0).float(),  # [3]
            'resize_mask': torch.from_numpy(resize_mask.astype(np.float32)).unsqueeze(dim=0),
            'resize_region': torch.from_numpy(resize_region*resize_mask).long(),
        }

    def __getitem__(self, index):
        item = self.all_lst[index]
        ds = None
        while ds is None:
            ds = self._load_data(item)
            item = random.choice(self.all_lst)

        if self.ori_image.any():
            ds['ori_img'] = self.ori_image
        return ds

    def __len__(self):
        return len(self.all_lst)

    @property
    def _error_data(self):
        cc = torch.tensor([0.])
        return {'img': cc, 'flag': False}

    def _load_model(self, model_path):
        if self.mode == 'eval':
            num_pt_mesh = self.num_pt_mesh_large
        else:
            num_pt_mesh = self.num_pt_mesh_small
        model_points = ply_vtx(model_path) / self.scale

        xmax, xmin = model_points[:, 0].max(), model_points[:, 0].min()
        ymax, ymin = model_points[:, 1].max(), model_points[:, 1].min()
        zmax, zmin = model_points[:, 2].max(), model_points[:, 2].min()
        extent = np.array([xmax-xmin, ymax-ymin, zmax-zmin])
        lf_border = np.array([xmin, ymin, zmin])
        dellist = [j for j in range(0, len(model_points))]
        dellist = random.sample(dellist, len(model_points) - num_pt_mesh)
        model_points = np.delete(model_points, dellist, axis=0)

        return model_points, extent, lf_border

    @staticmethod
    def rt2matrix(r, t):
        ext = np.eye(4, 4)
        ext[:3, :3] = r
        ext[:3, 3] = t
        return ext

border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
# border_list = [-1] + list(range(8, 680, 8)) + [680]

def read_lines(p):
    with open(p, 'r') as f:
        return [line.strip() for line in f.readlines()]

def ply_vtx(path):
    f = open(path)
    assert f.readline().strip() == "ply"
    f.readline()
    f.readline()
    N = int(f.readline().split()[-1])
    while f.readline().strip() != "end_header":
        continue
    pts = []
    for _ in range(N):
        pts.append(np.float32(f.readline().split()[:3]))
    return np.array(pts)

def get_bbox(bbox):
    bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
    if bbx[0] < 0:
        bbx[0] = 0
    if bbx[1] >= 480:
        bbx[1] = 479
    if bbx[2] < 0:
        bbx[2] = 0
    if bbx[3] >= 640:
        bbx[3] = 639
    rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]
    r_b = rmax - rmin
    for tt in range(len(border_list)):
        if border_list[tt] < r_b < border_list[tt + 1]:
            r_b = border_list[tt + 1]
            break
    c_b = cmax - cmin
    for tt in range(len(border_list)):
        if border_list[tt] < c_b < border_list[tt + 1]:
            c_b = border_list[tt + 1]
            break
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > 480:
        delt = rmax - 480
        rmax = 480
        rmin -= delt
    if cmax > 640:
        delt = cmax - 640
        cmax = 640
        cmin -= delt
    return rmin, rmax, cmin, cmax

def get_square_bbox(bbox, height_px=480, width_px=640):
    bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
    if bbx[0] < 0:
        bbx[0] = 0
    if bbx[1] >= 480:
        bbx[1] = 479
    if bbx[2] < 0:
        bbx[2] = 0
    if bbx[3] >= 640:
        bbx[3] = 639
    rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]

    rmax += 1
    cmax += 1
    r_b = rmax - rmin  # h
    c_b = cmax - cmin  # w
    if r_b <= c_b:
        r_b = c_b
    else:
        c_b = r_b
    for tt in range(len(border_list)):
        if border_list[tt] < r_b < border_list[tt + 1]:
            r_b = border_list[tt + 1]
            break

    for tt in range(len(border_list)):
        if border_list[tt] < c_b < border_list[tt + 1]:
            c_b = border_list[tt + 1]
            break

    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > height_px:
        delt = rmax - height_px
        rmax = height_px
        rmin -= delt
        if rmin < 0:
            rmax = rmax - rmin
            rmin = 0
            if rmax >= height_px:
                rmax = height_px - 1

    if cmax > width_px:
        delt = cmax - width_px
        cmax = width_px
        cmin -= delt
        if cmin < 0:
            cmax = cmax - cmin
            cmin = 0
            if cmax >= width_px:
                cmax = width_px - 1

    m = (rmax - rmin) - (cmax - cmin)
    if m > 0:
        rmax = rmax - np.floor(m / 2)
        rmin = rmin + np.floor(m / 2)
    elif m < 0:
        cmax = cmax + np.floor(m / 2)
        cmin = cmin - np.floor(m / 2)

    return int(rmin), int(rmax), int(cmin), int(cmax)

def main():
    root_path = '/root/Source/ymy_dataset/lm-bop'
    from numpy import array, float32
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt

    dataset = PoseDataset('train', 500, True, root_path, 0.000, 8, cls_type='ape')

    dataloader_test = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=0)
    ds_iter = dataloader_test.__iter__()
    for _ in range(5):
        datas = ds_iter.__next__()
        datas = {k:v for k,v in datas.items() if type(v).__name__ == 'Tensor'}
        print(datas['cls_id'].size())

def tst():
    # p = '/root/Source/ymy_dataset/lm-bop/train/000001/rgb'
    # paths = sorted(Path(p).glob('*.png'))
    # print(paths[0].stem)
    p = '/root/Source/ymy_dataset/lm-bop/train/000001/scene_gt.json'
    with open(p) as f:
        files = json.load(f)

    print(type(files))

if __name__ == '__main__':
    # tst()
    main()


import matplotlib.pyplot as plt
import numpy as np
import itertools
# from numba import njit
#import pandas as pd

#@njit
def transfer_matrix(N1, N2, polar='TE', n1=1, n2=1):
	tm = np.empty((2,2), dtype = np.complex)
	if polar == 'TE':
		tm[0, 0] = (N2 + N1) / (2. * N2)
		tm[0, 1] = (N2 - N1) / (2. * N2)
		tm[1, 0] = (N2 - N1) / (2. * N2)
		tm[1, 1] = (N2 + N1) / (2. * N2)
	else:
		tm[0, 0] = (n2 ** 2 * N1 + n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		tm[0, 1] = (n2 ** 2 * N1 - n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		tm[1, 0] = (n2 ** 2 * N1 - n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		tm[1, 1] = (n2 ** 2 * N1 + n1 ** 2 * N2) / (2. * n1 * n2 * N2)
	return tm

#@njit
def N_calculation(n_first_media, incident_angle, n_current_media):
	return np.sqrt(n_current_media ** 2 - n_first_media ** 2 * (np.sin(incident_angle)) ** 2, dtype=np.complex)

#@njit
def phase_shift(d, N, k0):
	p_shift = np.empty((2, 2), dtype = np.complex)
	p_shift[0, 0] = np.exp(1.j * d * N * k0)
	p_shift[0, 1] = 0.+0.j
	p_shift[1, 0] = 0.+0.j
	p_shift[1, 1] = np.exp(-1.j * d * N * k0)
	return p_shift

#@njit
def R_func(dictionary_structure={}, wl = 780, teta = 45, polar='TE'):
	list_N = []
	list_n = []
	T = np.eye(2, dtype=np.complex)
	tet = teta * np.pi / 180
	k0 = 2 * np.pi / wl
	n0 = complex(0., 0.)  # current n
	for i in range(0, len(dictionary_structure)):
		if i == 0:
			n0 = dictionary_structure[i]['n']
			list_N.append(N_calculation(n0, tet, n0))
			list_n.append(n0) # добавил n для TM
			continue
		if dictionary_structure[i]['name'] == 'ФК':
			n1 = dictionary_structure[i]['n1']
			n2 = dictionary_structure[i]['n2']
			d1 = dictionary_structure[i]['d1']
			d2 = dictionary_structure[i]['d2']
			num_layers = int(dictionary_structure[i]['N'])
			N1 = N_calculation(n0, tet, n1)
			N2 = N_calculation(n0, tet, n2)
			from_upper_to_1 = transfer_matrix(list_N[-1], N1, polar, list_n[-1], n1)
			from_1_to_2 = transfer_matrix(N1, N2, polar, n1, n2)
			from_2_to_1 = transfer_matrix(N2, N1, polar, n2, n1)
			F_n1 = phase_shift(d1, N1, k0) # x1 - d, x2 - N, x3 - k0
			F_n2 = phase_shift(d2, N2, k0)

			T = F_n2 @ from_1_to_2 @ F_n1 @ from_upper_to_1 @ T
			T_bilayer = F_n2 @from_1_to_2 @ F_n1 @from_2_to_1
			T = np.linalg.matrix_power(T_bilayer, num_layers - 1) @ T
			list_N.append(N1)
			list_N.append(N2)
			list_n.append(n1)
			list_n.append(n2)
		elif dictionary_structure[i]['name'] == 'Слой':
			n = dictionary_structure[i]['n']
			d = dictionary_structure[i]['d']
			N = N_calculation(n0, tet, n)
			from_this_to_next = transfer_matrix(list_N[-1], N, polar, list_n[-1], n)
			F = phase_shift(d, N, k0)
			T = np.dot(from_this_to_next, T)
			T = np.dot(F, T)
			list_N.append(N)
			list_n.append(n) # добавил n для TM
		elif dictionary_structure[i]['name'] == 'Среда':
			n = dictionary_structure[i]['n']
			N = N_calculation(n0, tet, n)
			from_last_to_sub = transfer_matrix(list_N[-1], N, polar, list_n[-1], n)
			T = np.dot(from_last_to_sub, T)
			list_N.append(N)
			list_n.append(n)
	return -(T[1, 0] / T[1, 1])

'''
get_rotation function take angle
in grad and return rotation matrix
'''
#@njit
def get_rotation(angle):
	rot_matrix = np.empty((2, 2), dtype = np.complex)
	rot_matrix[0, 0] = np.cos(np.pi * angle / 180)
	rot_matrix[0, 1] = np.sin(np.pi * angle / 180)
	rot_matrix[1, 0] = np.sin(np.pi * angle / 180)
	rot_matrix[1, 1] = -np.cos(np.pi * angle / 180)
	return rot_matrix

#@njit
def decor1(angles, wl, struct, polar):
	return R_func(struct, wl, angles, polar)

vectorize_R = np.vectorize(decor1)
vectorize_N = np.vectorize(N_calculation)

#@njit
def transfer_matrix_vec(N1, N2, polar='TE', n1=1+0.j, n2=1+0.j):
	if polar == 'TE':
		t00 = (N2 + N1) / (2. * N2)
		t01 = (N2 - N1) / (2. * N2)
		t10 = (N2 - N1) / (2. * N2)
		t11 = (N2 + N1) / (2. * N2)
	else:
		t00 = (n2 ** 2 * N1 + n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		t01 = (n2 ** 2 * N1 - n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		t10 = (n2 ** 2 * N1 - n1 ** 2 * N2) / (2. * n1 * n2 * N2)
		t11 = (n2 ** 2 * N1 + n1 ** 2 * N2) / (2. * n1 * n2 * N2)
	return np.vstack((t00, t01, t10, t11)).reshape((2, 2, -1)).T

#@njit
def phase_shift_vec(d, N, k0):
	p00 = np.exp(1.j * d * N * k0)
	p01 = np.zeros(np.shape(p00), dtype=np.complex)
	p10 = np.zeros(np.shape(p00), dtype=np.complex)
	p11 = np.exp(-1.j * d * N * k0)
	return np.vstack((p00, p01, p10, p11)).reshape((2, 2, -1)).T

#@njit
def generate_TMM(n0, n1, n2, angles, polar):
	N1 = N_calculation(n0, np.pi / 180 * angles, n1)
	N2 = N_calculation(n0, np.pi / 180 * angles, n2)
	TMMs = transfer_matrix_vec(N1, N2, polar, n1, n2)
	return TMMs

#@njit
def generate_TMM_to_metal(n0, n2, n3, angles, polar, zk, TMM1):
	N2 = N_calculation(n0, np.pi / 180 * angles, n2)
	N3 = N_calculation(n0, np.pi / 180 * angles, n3)
	TMM2 = transfer_matrix_vec(N2, N3, polar, n2, n3)
	Phases = phase_shift_vec(zk, N2, 1.)
	TMMs = TMM2 @ Phases @ TMM1
	return TMMs

#@njit
def get_K_xz_air(k_mkm, k_mkm_gap, angles):
	k_x_b = k_mkm * np.sin(np.pi / 180 * angles, dtype=np.complex)
	k_z_b = np.sqrt(k_mkm_gap ** 2 - k_x_b ** 2, dtype=np.complex)
	k_matrix = np.vstack((k_x_b, k_z_b))
	return k_matrix

#@njit
def E_air_spectrum(s, R, TMMs_):
	E_air = np.empty((2, np.size(s)), dtype=np.complex)
	E_inc = np.array(np.hstack((s, s * R)).T, dtype=np.complex)

	# tmp = E_inc.T.reshape((np.size(s), 2, 1))
	# res = TMMs_ @ tmp
	# E_plus = res[:, 0, 0].reshape((-1, 1))
	# E_minus = res[:, 1, 0].reshape((-1, 1))

	for i in range(np.size(s)):
		E_air[:, i] = TMMs_[i] @ E_inc[:, i]
	E_plus = E_air[0].reshape((-1, 1))
	E_minus = E_air[1].reshape((-1, 1))
	return E_plus, E_minus

#@njit
def get_defocus_phase(k_mkm, angles, L):
	if L == 0:
		return 1.
	median_angle = np.pi / 180 * angles[np.size(angles) // 2]
	h_shift = L * np.cos(median_angle)
	x_shift = L * np.sin(median_angle)
	k_x = k_mkm * np.sin(np.pi / 180 * angles)
	k_z = k_mkm * np.cos(np.pi / 180 * angles)
	phase_shift = k_x * x_shift + k_z * h_shift
	return np.exp(1.0j * phase_shift).reshape((-1, 1))

#@njit
def spectral(a, k):
	if a == np.inf:
		return np.array([1.], dtype=np.float64)
	return np.sqrt(a)/(2 * np.sqrt(np.pi)) * np.exp(-a**2 * k**2 / 4)

#@njit
def get_k_x(a=5, N_x=401, decr_times=20):
	if a == np.inf:
		return np.array([0.], dtype=np.float64)
	k_min = 2 * np.sqrt(np.log(decr_times)) / a
	dk_x = 2 * k_min / N_x
	return np.arange(-N_x // 2, N_x // 2) * dk_x

def cut_BSW_intesity(angles_range, X, struct, a=5, Z=0.3, L=0, polar='TE', wl=780., plasmon=False):
	'''
	set default parameters range
	'''
	# _, ax = plt.subplots(1, 4, figsize=(20, 5))

	k = 2 * np.pi * 1.0e9 / wl * struct[0]['n']
	k_mkm = k / 1.0e6
	k_mkm_gap = k_mkm / struct[0]['n'] * struct[1]['n']
	k_x = get_k_x(a=a, N_x=401, decr_times=5)
	dk = k_x[1] - k_x[0]
	s = spectral(a, k_x).reshape((-1, 1))
	r_pos = np.empty((1, 2))
	r_pos[0,1] = Z
	E_destr = np.empty((np.size(angles_range), np.size(X)), dtype=np.complex)
	for j, alpha in enumerate(angles_range):
		angles = 180 / np.pi * np.arcsin(k_x / k_mkm) + alpha
		R = vectorize_R(angles, wl, struct, polar).reshape((-1, 1))
		TMM1 = generate_TMM(struct[0]['n'], struct[0]['n'], struct[1]['n'], angles, polar)
		if plasmon == True:
			TMM2 = generate_TMM_to_metal(struct[0]['n'], struct[1]['n'], struct[2]['n'],
										angles, polar, Z * k_mkm / struct[0]['n'], TMM1)
			TMMs = generate_TMM_to_metal(struct[0]['n'], struct[2]['n'], struct[3]['n1'],
										angles, polar, struct[2]['d'] / 1000. * k_mkm / struct[0]['n'], TMM2)
			r_pos[0,1] = 0
		else:
			TMMs = TMM1
		k_xz = get_K_xz_air(k_mkm, k_mkm_gap, angles)
		defocus_phase = get_defocus_phase(k_mkm, angles, L) # эта фаза считается в призме, поэтому вычтсленные выше k_xz для воздуха не используются
		E_plus, E_minus = E_air_spectrum(s, R, TMMs)
		E_forward = np.zeros(np.size(X), dtype=np.complex)
		E_backward = np.zeros(np.size(X), dtype=np.complex)

		# if j % 8 == 0:
		# 	ax[0].plot(angles.real, abs(E_plus))
		# 	ax[1].plot(angles.real, abs(E_minus))
		# 	ax[2].plot(angles.real, abs(R) + 0.05 * (j // 8))
		# 	ax[3].plot(angles.real, np.angle(R) + 0.05 * (j // 8))

		# if angles[0].real < 42.2:
		# 	E_destr[j] = np.zeros(np.size(X), dtype=np.complex)
		# 	continue

		# coord_arr = np.vstack((X, Z * np.ones(np.size(X)))).T
		# E_forward = np.exp(1.0j * (coord_arr @ k_xz), dtype=np.complex) @ (E_plus * defocus_phase) * dk
		# coord_arr = np.vstack((X, -Z * np.ones(np.size(X)))).T
		# E_backward = np.exp(1.0j * (coord_arr @ k_xz), dtype=np.complex) @ (E_minus * defocus_phase) * dk
		# E_destr[j] = (E_forward + E_backward).T

		for i, x in enumerate(X):
			r_pos[0, 0] = x
			E_forward[i] = np.exp(1.0j * r_pos @ k_xz, dtype=np.complex) @ (E_plus * defocus_phase) * dk
			r_pos[0, 1] *= -1
			E_backward[i] = np.exp(1.0j * r_pos @ k_xz, dtype=np.complex) @(E_minus * defocus_phase) * dk
			r_pos[0, 1] *= -1
		E_destr[j] = E_forward + E_backward
	return E_destr

#@njit
def plot_2Dmap(data, extent=None, isE=False):
	plt.figure(figsize=(9, 5))
	if isE == True:
		proc_data = data
	else:
		proc_data = abs(data) ** 2
	plt.imshow(proc_data, cmap=plt.cm.hot, origin='lower', aspect='auto', extent=extent)
	cax = plt.axes([0.95, 0.13, 0.05, 0.75])
	plt.colorbar(cax=cax)
	plt.show()

#@njit
def E_beam_calc(X, Z, s, k_matrix, mode='i', alpha=0, l_0=0, plot_I=True):
	'''
	This function compute field distribution by matrix multiplications.
	grid - array of coordinates in form:
		[(x_0, z_0), (x_0, z_1), ..., (x_0, z_n),
						...,
		(x_m, z_0), (x_m, z_1), ..., (x_m, z_n)]

	'''
	rot_matrix = get_rotation(alpha)
	dk = k_matrix[0][1] - k_matrix[0][0]
	h_shift = 0. if l_0 == 0 else l_0 * np.cos(np.pi / 180 * alpha)
	sign = 1. if mode == 'i' else -1
	grid = np.array(list(itertools.product(X, sign * Z - h_shift)), dtype=np.complex)
	E = np.exp(1.0j * (grid @ rot_matrix.T) @ k_matrix, dtype=np.complex) @ s * dk
	E_res = E.reshape((np.size(X), -1)).T
	if plot_I == True:
		plot_2Dmap(E_res)
	return E_res

#@njit
def vizualize_beam(X, Z, struct, alpha, l_0=0, a=5, mode='r', wl=780., polar='TE', plot=False):

	'''
	This function return complex value E (electromagnetic field)
	distribution in the modulation field and plot Intensity
	distribution, that can be set throught input parameters.

	Parameters:
		X - x-range coorinates
		Z - z-range coorinates
		struct - structure (Otto configuration for example) description
				in dictionaty form
		alpha - incidence angle (degree)
		l_0 - defocusing length (mkm)
		mode - simalation regime (r - reflected, i - incidence, f - full)
		wl - wavelength (nm)
	'''

	k = 2 * np.pi * 1.0e9 / wl * struct[0]['n']
	k_mkm = k / 1.0e6
	k_x = get_k_x(a=a, N_x=401)
	k_z = np.sqrt(k_mkm ** 2 - k_x ** 2, dtype=np.complex)
	s = spectral(a=a, k=k_x).reshape((-1, 1))
	k_matrix = np.vstack((k_x, k_z))
	s_angle_range = 180 / np.pi * np.arcsin(k_x / k_mkm)

	angles = s_angle_range + alpha
	R = vectorize_R(angles, wl, struct, polar).reshape((-1, 1))
	if mode == 'r' or mode == 'f':
		E_refl = E_beam_calc(X, Z, s * R, k_matrix, 'r', alpha, l_0, plot)
	if mode == 'i' or mode == 'f':
		E_inc = E_beam_calc(X, Z, s, k_matrix, 'i', alpha, l_0, plot)

	if mode == 'f':
		return [E_inc, E_refl, E_inc + E_refl]
	elif mode == 'r':
		return [E_refl]
	elif mode == 'i':
		return [E_inc]

def get_max_angle(E_bsw, return_shift=None, plot_I_a=None, plot_I_x=None):
    I = np.abs(E_bsw) ** 2
    I_max = []
    for line in I:
        I_max.append(np.max(line))
    if plot_I_a != None:
        plt.plot(angles_range, I_max)
        plt.grid()
        plt.xlabel("Inc angle")
        plt.ylabel("Intensity")
    angle_max = angles_range[np.argmax(I_max)]
    if return_shift != None:
        if plot_I_x != None:
            plt.plot(X_range, I[np.argmax(I_max)])
        shift = X_range[np.argmax(I[np.argmax(I_max)])]
        return [angle_max, shift, np.max(I)]
    print("I_max: ", np.max(I_max), "angle_max: ", angles_range[np.argmax(I_max)])
    return [angle_max]


import sys
sys.stdin = open('input (1).txt')

for tc in range(1, 11):
    num = int(input())
    arr = [list(map(int, input().split())) for _ in range(100)]
    result = 0

    for j in range(100):
        for i in range(100):
            if arr[i][j] == 1:                      # N극일 때
                if i+1 < 100 and arr[i+1][j] == 0:  # 아랫줄이 범위 내고 0이라면
                    arr[i][j], arr[i+1][j] = arr[i+1][j], arr[i][j] # 아랫줄의 값과 바꾸기(1을 한 칸 내리기)
                if i == 99:             # 만약 i가 99고 해당 자리가 1이면
                    arr[i][j] = 0       # 밖으로 나가게 되므로 해당 자리를 0으로 바꾸기
                else:
                    continue

            if arr[i][j] == 2:                      # S극일 때도 윗줄에 대해 똑같이 적용
                if 0 <= i-1 and arr[i-1][j] == 0:
                    arr[i][j], arr[i-1][j] = arr[i-1][j], arr[i][j]
                if i == 0:
                    arr[i][j] = 0
                else:
                    continue

    for i in range(100):
        for j in range(100):
            if arr[i][j] == 1:          # N극에 대해
                if arr[i+1][j] == 2:    # 아랫줄이 S극이면 교착상태이므로 result += 1
                    result += 1
    # N극은 아래로만 내려가고, N+S극 조합일 때 교착상태이므로 하나에 대해서만 조사사
    print(f'#{tc} {result}')


"""Name hold a name choice for a Request
"""
# from . import db, ma
from marshmallow import fields
from sqlalchemy import event
from sqlalchemy.orm import backref
from sqlalchemy.orm.attributes import get_history

from namex.models import db, ma

class Name(db.Model):
    __tablename__ = 'names'

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(1024), index=True)
    state = db.Column(db.String(15), default='NE')  # NE=Not Examined; R=Rejected; A=Accepted; C=Cond. Accepted
    choice = db.Column(db.Integer)
    designation = db.Column(db.String(50), default=None)
    consumptionDate = db.Column('consumption_date', db.DateTime(timezone=True))
    corpNum = db.Column('corp_num', db.String(10), default=None)
    remoteNameId = db.Column('remote_name_id', db.BigInteger)

    # Decision info
    conflict1 = db.Column(db.String(250), default='')  # optional conflict name
    conflict2 = db.Column(db.String(250), default='')  # optional conflict name
    conflict3 = db.Column(db.String(250), default='')  # optional conflict name
    conflict1_num = db.Column(db.String(250), default='')  # optional conflict name - corp or NR number
    conflict2_num = db.Column(db.String(250), default='')  # optional conflict name - corp or NR number
    conflict3_num = db.Column(db.String(250), default='')  # optional conflict name - corp or NR number
    decision_text = db.Column(db.String(1000), default='')

    nrId = db.Column('nr_id', db.Integer, db.ForeignKey('requests.id'), index=True)
    commentId = db.Column('comment_id', db.Integer, db.ForeignKey('comments.id'))
    # nameRequest = db.relationship('Request')

    # if a comment is added during decision, link it to the name record to be sent back to NRO
    comment = db.relationship("Comment", backref=backref("related_name", uselist=False), foreign_keys=[commentId])

    # Required for name request name analysis
    _name_type_cd = db.Column('name_type_cd', db.String(10))

    NOT_EXAMINED = 'NE'
    APPROVED = 'APPROVED'
    REJECTED = 'REJECTED'
    CONDITION = 'CONDITION'
    # Needed for name request reservation before completing the nr
    RESERVED = 'RESERVED'
    COND_RESERVE = 'COND-RESERVE'

    # Properties added for Name Request
    @property
    def name_type_cd(self):
        """Property containing the name type which is used by name Request."""
        return self._name_type_cd

    @name_type_cd.setter
    def name_type_cd(self, value: str):
        self._name_type_cd = value

    def as_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'name_type_cd': self.name_type_cd,
            'designation': self.designation,
            'choice': self.choice,
            'state': self.state,
            'conflict1': self.conflict1,
            'conflict2': self.conflict2,
            'conflict3': self.conflict3,
            'conflict1_num': self.conflict1_num,
            'conflict2_num': self.conflict2_num,
            'conflict3_num': self.conflict3_num,
            'decision_text': self.decision_text,
            'consumptionDate': self.consumptionDate.isoformat() if self.consumptionDate else None,
            'corpNum': self.corpNum,
            'comment': None if self.comment is None else self.comment.as_dict()
        }

    @classmethod
    def find_by_name(cls, name):
        return cls.query.filter_by(name=name).first()

    def save_to_db(self):
        # force uppercase names
        self.name = self.name.upper()

        db.session.add(self)
        db.session.commit()

    def add_to_db(self):
        db.session.add(self)

    def delete_from_db(self):
        db.session.delete(self)
        db.session.commit()

@event.listens_for(Name, 'after_insert')
@event.listens_for(Name, 'after_update')
def update_nr_name_search(mapper, connection, target):
    """Add any changes to the name to the request.nameSearch column and publish name state changes where applicable."""
    from flask.globals import current_app

    from namex.models import Event, Request, State
    from namex.services.audit_trail.event_recorder import EventRecorder

    name = target
    nr = Request.find_by_id(name.nrId)
    if nr:
        # get the names associated with the NR
        names_q = connection.execute(
            f"""
            SELECT names.name from names
            JOIN requests on requests.id = names.nr_id
            WHERE requests.id={nr.id}
            """
        )
        # format the names into a string like: |1<name1>|2<name2>|3<name3>
        names = [x[0] for x in names_q.all()]
        name_search = ''
        for item, index in zip(names, range(len(names))):
            name_search += f'|{index + 1}{item}{index + 1}|'
        # update the name_search field of the nr with the formatted string
        connection.execute(
            """
            UPDATE requests
            SET name_search=%s
            WHERE id=%s
            """,
            ('(' + name_search + ')', nr.id)
        )

        # set nr state to consumed
        name_consume_history = get_history(name, 'consumptionDate')
        current_app.logger\
            .debug('name_consume_history check - nrNum: {}, consumptionDate: {}, corpNum: {}, state: {}'
            .format(nr.nrNum, name.consumptionDate, name.corpNum, name.state))
        # Note: we cannot just check for a corpNum addition due to some Society change of name NRs coming over from
        #  NRO extractor providing a value for the corpNum field.
        if len(name_consume_history.added) \
                and name.consumptionDate \
                and name.corpNum \
                and name.state in ['APPROVED', 'CONDITION']:
            # Adding an after_flush_postexec to avoid connection and transaction closed issue's
            # Creating one time execution event when ever corpNum is added to a name
            # corpNum sets from nro-extractor job
            @event.listens_for(db.session, 'after_flush_postexec', once=True)
            def receive_after_flush_postexec(session, flush_context):
                nr = Request.find_by_id(name.nrId)
                nr.stateCd = State.CONSUMED
                nr.add_to_db()
                current_app.logger.debug('moved to CONSUMED state {}'.format(name.corpNum))
                EventRecorder.record_as_system(Event.UPDATE_FROM_NRO, nr, {
                    'id': nr.id,
                    'nrNum': nr.nrNum,
                    'stateCd': nr.stateCd
                }, True)
                current_app.logger.debug('moved to CONSUMED state event logged {}'.format(nr.nrNum))

class NameSchema(ma.SQLAlchemySchema):
    class Meta:
        model = Name
        fields = (
            'choice',
            'comment',
            'conflict1',
            'conflict2',
            'conflict3',
            'conflict1_num',
            'conflict2_num',
            'conflict3_num',
            'consumptionDate',
            'corpNum',
            'decision_text',
            'designation',
            'id',
            'name_type_cd',
            'name',
            'state'
        )

    conflict1 = fields.String(required=False, allow_none=True)
    conflict2 = fields.String(required=False, allow_none=True)
    conflict3 = fields.String(required=False, allow_none=True)
    conflict1_num = fields.Field(required=False, allow_none=True)
    conflict2_num = fields.Field(required=False, allow_none=True)
    conflict3_num = fields.Field(required=False, allow_none=True)
    decision_text = fields.String(required=False, allow_none=True)
    comment = fields.String(required=False, allow_none=True)
    consumptionDate = fields.DateTime(required=False, allow_none=True)
    corpNum = fields.String(required=False, allow_none=True)
    designation = fields.String(required=False, allow_none=True)
    name = fields.String(
        required=True,
        error_messages={'required': {'message': 'name is a required field'}}
    )
    name_type_cd = fields.String(required=False, allow_none=True)


import optparse

parser = optparse.OptionParser(
    usage='%prog COMMAND [OPTIONS]',
    version="x.x.x",
    add_help_option=False)

parser.add_option(
    '-h', '--help',
    dest='help',
    action='store_true',
    help='Show help')

parser.disable_interspersed_args()


import ncs

is_lsa = False
ietf_l2vpn_servicepoint = "ietf-l2vpn-ntw-servicepoint"
ietf_l2vpn_validation_callpoint = "ietf-l2vpn-nm-validation"
l2vpn_ntw_augmentations_y1731_servicepoint = "l2vpn-ntw-augmentations-y1731-servicepoint"

def is_lsa_setup():
    with ncs.maapi.single_read_trans("", "system", db=ncs.RUNNING) as th:
        global is_lsa
        try:
            root = ncs.maagic.get_root(th)
            # If this internal service path exists, it is not LSA
            non_lsa_path = root.\
                cisco_flat_L2vpn_fp_internal_local_site__flat_L2vpn_internal_local_site_service
            if non_lsa_path:
                is_lsa = False
        except Exception:
            is_lsa = True

def get_evpn_pbb_comp_name(vpn_node_id, vpn_network_access_id):
    return f"{vpn_node_id}_{vpn_network_access_id}"


import numpy as np # Bring numpy in for indexing via mask

d = ["able", "ale", "hullabuloo", "apple"] # Here is our dictionary
subsq = [] # emtpy array to hold all of our subsequences

def yieldSubSq(s_, d):
    '''
    arguments:
    s_ string: Master sequence containig subsequences
    d list of strings: dictionary of words to check
    '''
    for w in d: # For each word in the dictionary
        s = s_ # let s be our master string
        print(f"Word: {w}")
        isSubSq=True # Assume all words are subsequences until proven otherwise
        for l in w: # for each leter in the word w
            isInSq=False # keep track of whether the letter l is in our master sequence
            for i in range(len(s)-1): # for each letter i the master sequence
                if(l==s[i]): # if the letter l is in the sequence
                    s=s[i+1:] # remove every letter in master sequence up to the current letter
                    isInSq=True # letter l is in the sequence
                    break
            if(isInSq==False): # if letter l wasn't in the sequence
                isSubSq=False # the word is not a subsequence
                break
        if(isSubSq): # if the word is a subsequence
            yield w # yield it

s = "abppplee" # define the master sequence
subSqGen = yieldSubSq(s, d)  # declare the generator of subsequences

for subSq in subSqGen:
    subsq.append(subSq) # yield all available subsequences in dictionary

subsq=np.array(subsq) # convert to numpy array for boolean indexing
lengths = np.array(list(map(len, subsq))) # Create array of lengths of each subsequence
print(subsq[np.where(lengths==np.max(lengths))]) # print out the longest subsequence


#!pipenv run python3
from pprint import pprint
from PyInquirer import prompt, Separator
from git import Repo
import os
import sys

def get_git_root():
    try:
        git_repo = Repo(os.getcwd(), search_parent_directories=True)
        git_root = git_repo.git.rev_parse("--show-toplevel")
        return git_root
    except:
        print("Not currently in a git repo")
        sys.exit()

repo = Repo(get_git_root())

assert not repo.bare

questions = [
    {
        'type': 'checkbox',
        # 'qmark': 'ðŸ˜ƒ',
        'message': 'Select Branches to Delete',
        'name': 'branches',
        'choices': [
            {
                'name': f"{b} -> {b.tracking_branch()}",
                'value': b
            }
            for b in repo.branches],
        'validate': lambda answer: 'You must choose at least one topping.' \
            if len(answer) == 0 else True
    }
]

answers = prompt(questions)
print("Deleted branches:")
try:
    for b in answers['branches']:
        b.delete(repo, b)
        print(b)
except KeyError:
    pass


'''
A binary watch has 4 LEDs on the top which represent the hours (0-11), and the 6 LEDs on the bottom represent
the minutes (0-59).

Given a non-negative integer n which represents the number of LEDs that are currently on, return all
possible times the watch could represent.
'''

def readBinaryWatch(n):
	'''
	This function goes through each time, decides how many LEDs are on, and if
	the number of LEDs for a time matches the target number of LEDs n, then we append to our output of
	possible times with n LEDs.
	'''
	output = []

	for h in range(0,12):
		for m in range(0,60):
			# count the LEDs that are on in the hours part by looking at the 1's in its binary representation
			hLEDS = bin(h).count('1')
			# count the LEDs that are on in the minutes part by looking at the 1's in its binary representation
			mLEDS = bin(m).count('1')
			if mLEDS + hLEDS == n:
				output.append('{hour}:{minute:02}'.format(hour=h, minute=m))
	return output

print(readBinaryWatch(3))


""" Enhancement of pathlib, representation with a tree data-structure.

This module enhances the functionalities of pathlib, interpreting the Path
objects as nodes of a tree and automatically creating their subtrees. The tree
can be explored for analysis purposes.
"""

from   pathlib import Path
from   typing import Union, Callable, Iterator, Any
from   enum import Enum
import copy
import xlsxwriter

__all__ = ["PathNode", "PathTree", "PathTreeProperty", "TreeException", "Size"]

class PathTreeProperty(Enum):
    """ Enumerator class for the tree properties.
    """

    HEIGHT = "<height>"
    """ Height (of a node): the minimal distance from the note to a leaf. """
    DEPTH = "<depth>"
    """ Depth (of a node): the distance from the note to the root. """
    NUM_OF_DIRECTORIES = "<num_dir>"
    """ No. of Directory (of a subtree): the number of directory in the subtree. """
    NUM_OF_FILES = "<num_file>"
    """ No. of Files (of a subtree): the number of files in the subtree. """
    NUM_OF_INODES = "<num_inode>"
    """ No. of inner nodes (of a subtree): the number of inner nodes in the subtree. """
    NUM_OF_LEAVES = "<num_leaves>"
    """ No. of leaves (of a subtree): the number of leaves in the subtree. """
    NUM_OF_NODES = "<num_nodes>"
    """ No. of nodes (of a subtree): the number of nodes in a subtree. """
    SIZE = "<size>"
    """ Size (of a subtree): the sum of the byte occupied by all the files in the subtree. """
    SIMPLE_SIZE = "<simple_size>"
    """ Size (of a subtree) in simplified format: the sum of the byte occupied by all the files in
    the subtree, expressed in KB, ..., TB accordingly. """
    PRUNED = "<pruned>"
    """ Pruned status (of a node): true if the node has been virtually removed from the tree, false
    otherwise. """

class Size(Enum):
    """ Enumerator class for the main byte multiples.
    """

    KB = 1024 ** 1
    """ Number of bytes in a kilobyte. """
    MB = 1024 ** 2
    """ Number of bytes in a megabyte. """
    GB = 1024 ** 3
    """ Number of bytes in a gigabyte. """
    TB = 1024 ** 4
    """ Number of bytes in a terabyte. """

    @staticmethod
    def simplified_size(size:int):
        """ Convert num of bytes to a sinthetic string.

        The size expressed with a multiple of the byte and with 1-3 significant
        digits. E.g. 23 MB
        """

        return \
            f"{size} B" if size < Size.KB.value else (
            f"{size // Size.KB.value} KB" if size < Size.MB.value else (
            f"{size // Size.MB.value} MB" if size < Size.GB.value else (
            f"{size // Size.GB.value} GB" if size < Size.TB.value else
            f"{size // Size.TB.value} TB"
            )))

class TreeException(Exception):
    """ Tree structure violated Exception.
    """

class PathNode():
    """ The PathNode class describe a tree node containing information
    associated to a path.

    The children of a node are sorted by: type (directories first, then files),
    path (alphabetically on the file/dir name).

    Attributes:
        path: the path associated to the node
        parent: the parent directory of the node, as a :class:`PathNode` object
        children: the children files/directories of the node, as list of\
        :class:`PathNode` objects
        data: a dictionary that can contain custom properties of the node\
        (e.g. height, subtree size, ...)
    """

    def __init__(self, path:Path, parent:Union["PathNode",None]=None) -> None:
        """ Create a new path node and all its subtree.

        Parameters:
            path: the path of the node
            parent: the path of the parent node (if None, it is a root node)
        """

        self.path = path
        self.parent = parent
        self.children = self.__compute_children()
        self.property = {}

    def __compute_children(self):
        """ Iteratively scan and generate the subtree of the node.
        """

        children = []
        try:
            if self.path.is_dir():
                for child in self.path.iterdir():
                    children.append(PathNode(child, self))
                children.sort(key=lambda node:f"{0 if node.path.is_dir() else 1}|{node.path.name}")
        except PermissionError:
            pass
        except FileNotFoundError:
            pass
        return children

    def is_root(self) -> bool:
        """ Return true if and only if the node is root.

        Check if the node has a parent node.

        Returns:
            True if and only if the node is root.
        """

        return self.parent is None

    def is_leaf(self) -> bool:
        """ Return true if and only if the node is a leaf.

        Check if the node has children nodes.

        Returns:
            True if and only if the node is a leaf.
        """

        return len(self.children) == 0

    def is_inode(self) -> bool:
        """ Return true if and only if the node is an inner node.

        Check if the node has children nodes.

        Returns:
            True if and only if the node is an inner node.
        """

        return not self.is_leaf()

    def is_file(self) -> bool:
        """ Return true if and only if the node is a file.

        Check if the node is a file.

        Returns:
            True if and only if the node is a file.
        """

        return not self.is_dir()

    def is_dir(self) -> bool:
        """ Return true if and only if the node is a directory.

        Check if the node is a directory.

        Returns:
            True if and only if the node is a directory.
        """

        return self.path.is_dir()

    def compute_bottom_up_property(
            self,
            property_name:Union[str,PathTreeProperty],
            base_func:Callable[["PathNode"], Any],
            recursive_func:Callable[["PathNode", list["PathNode"]], Any]
            ) -> None:
        """ Compute a property of bottom-up type in the subtree of the node.

        The bottom-up properties are recursively computed from the leaves
        (files) to the root: when the node is a leaf the proerty can be computed
        directly (without involving other nodes), when the node is an inner node
        the property is computed based also on the children nodes properties.
        Hence there are 2 functions that must be specified: the base case
        function (for leaves), the recursive function (for inner nodes).

        Parameters:
            property_name: the name of the property (used as key in property\
            dict)
            base_func: the function to compute the property on the leaves
            recursive_func: the function to compute the property on the inner\
            nodes (assuming it is already computed on the leaves)

        Parameters of base_func:

        * leaf (PathNode): the current node

        Parameters of recursive_func:

        * inode (PathNode): the current node
        * children (list(PathNode)): the list of children of the current node

        Examples::

            Compute the height.

            >>> base_func = lambda node: 0
            >>> recursive_func = lambda node, children: 1 + min([
            >>>     int(child.property["height"])
            >>>     for child in children
            >>> ])
            >>> root.compute_bottom_up_property("height", base_func, recursive_func)
        """

        if self.is_leaf():
            self.property[property_name] = base_func(self)
        else:
            for child in self.children:
                child.compute_bottom_up_property(property_name, base_func, recursive_func)
            self.property[property_name] = recursive_func(self, self.children)

    def compute_top_down_property(
            self,
            property_name:Union[str,PathTreeProperty],
            root_func:Callable[["PathNode"], Any],
            parent_func:Callable[["PathNode", "PathNode"], Any]
            ) -> None:
        """ Compute a property of top-down type in the subtree of the node.

        The top-down properties are recursively computed from the root to the
        leaves: when the node is root the proerty can be computed directly
        (without involving other nodes), when the node is a not-root node the
        property is computed based also on the parent node properties. Hence
        there are 2 functions that must be specified: the root function (for the
        root), the parent function (for not-root nodes).

        Parameters:
            property_name: the name of the property (used as key in property\
            dict)
            root_func: the function to compute the property on the root
            parent_func: the function to compute the property on the not-root\
            nodes (assuming it is already computed on the parent)

        Parameters of root_func:

        * root (PathNode): the current node

        Parameters of parent_func:

        * node (PathNode): the current node
        * parent (PathNode): the parent of the current node

        Examples::

            Compute the depth.

            >>> root_func = lambda root: 0
            >>> parent_func = lambda node, parent: 1 + parent.property["depth"]
            >>> root.compute_top_down_property("depth", root_func, parent_func)
        """

        if self.is_root():
            self.property[property_name] = root_func(self)
        else:
            if self.parent is not None:
                self.property[property_name] = parent_func(self, self.parent)
            else:
                raise TreeException(f"Not-root node has not a parent ({self})")
        for child in self.children:
            child.compute_top_down_property(property_name, root_func, parent_func)

    def compute_individual_property(
            self,
            property_name:Union[str,PathTreeProperty],
            property_func:Callable[["PathNode"], Any]
        ):
        """ Compute an individual property in the subtree of the node.

        The property does not exploit any knowledge on other nodes. However can
        exploit information stored in the node property dict (e.g. previously
        computed height).

        Parameters:
            property_name: the name of the property (used as key in property \
            dict)
            property_func: the function to compute the property on the node

        Parameters of property_func:

        * node (PathNode): the current node

        Examples::

            Compute a flag to identify directories.

            >>> property_func = lambda node: node.path.is_dir()
            >>> root.compute_individual_property("is_dir", property_func)
        """

        self.property[property_name] = property_func(self)
        for child in self.children:
            child.compute_individual_property(property_name, property_func)

    def remove_property(self, property_name:Union[str,PathTreeProperty]) -> bool:
        """ Remove a property of the node.

        Parameters:
            property_name: the name of the property to remove.

        Returns:
            True if the property is successfully removed, false if the property\
            is missing.
        """

        if property_name in self.property:
            self.property.pop(property_name)
            return True
        return False

    def copy(self) -> "PathNode":
        """ Return a deepcopy of the node.

        The pointers to other nodes are copied without modification hence point
        to the same nodes.

        Returns:
            A copy of the node.
        """

        return copy.deepcopy(self)

    def __str__(self) -> str:
        """ Return a string describing the node properties.

        Print the node name then the properties.

        Returns:
            A string describing the node.
        """

        properties = []
        for prop, value in self.property.items():
            if isinstance(prop, PathTreeProperty):
                properties.append(f"{prop.name}={value}")
            else:
                properties.append(f"{prop}={value}")
        return f"[Pathnode] > Path: '{self.path}' | Properties: {', '.join(properties)}"

class PathTree():
    """ The PathTree class describe a tree made up by PathNode nodes.

    The structure mimic the directory tree but add analytic functionalities.

    Attributes:

        root: the root PathNode of the tree
        property: the properties of the tree (equivalent to the root properties)
    """

    def __init__(self, root:Union[str, Path, PathNode], skip_properties:bool=False) -> None:
        """ Create a new PathTree, based on the root node.

        Parameters:
            root: the root path
            skip_properties: ignore computation of tree's main properties.
        """

        if isinstance(root, str):
            root = Path(root)
        if isinstance(root, Path):
            root = PathNode(root)
        self.root = root
        self.property = root.property
        if not skip_properties:
            self.__compute_basic_properties()

    def __iter__(self) -> Iterator:
        """ Return a default iterator on the nodes of the tree.

        The order is breadth-first and uses the iterator returned by the
        `breadth_first_iter()` function. To use the depth/first order, use the
        iterator returned by the `depth_first_iter()` function.

        Returns:
            An iterator for the nodes of the tree.
        """

        return self.breadth_first_iter()

    def breadth_first_iter(self) -> Iterator:
        """ Return an iterator on the nodes of the tree using breadth-first
        order.

        Returns:
            An iterator for the nodes of the tree, in breadth-first order.
        """

        nodes = [self.root]
        while len(nodes) > 0:
            node = nodes.pop(0)
            nodes = nodes + node.children
            yield node

    def depth_first_iter(self) -> Iterator:
        """ Return an iterator on the nodes of the tree using depth-first
        order.

        Returns:
            An iterator for the nodes of the tree, in depth-first order.
        """

        nodes = [self.root]
        while len(nodes) > 0:
            node = nodes.pop(0)
            nodes = node.children + nodes
            yield node

    def validated_iter(self, valid_func:Callable[[PathNode], bool]) -> Iterator:
        """ Return an iterator on filtered nodes of the tree.

        The nodes that do not satisfy the condition are excluded and so their
        subtree.

        Parameters:
            valid_func: the criteria to keep nodes.

        Returns:
            An iterator that exclude not valid nodes and subtrees.

        Parameters of valid_func:

        * node (PathNode): the node to test

        Return of valid_func:

        * True if the node is acceptable, false otherwise.
        """

        nodes = [self.root] if valid_func(self.root) else []
        while len(nodes) > 0:
            node = nodes.pop(0)
            nodes = [child for child in node.children if valid_func(child)] + nodes
            yield node

    def compute_bottom_up_property(
            self,
            property_name:Union[str,PathTreeProperty],
            leaf_func:Callable[["PathNode"], Any],
            inode_func:Callable[["PathNode", list["PathNode"]], Any]
            ) -> None:
        """ Compute a property of bottom-up type in the tree.

        This function calls the namesake function of the `PathNode` on the root.
        Check `PathNode.compute_bottom_up_property()` for more details.

        Parameters:
            property_name: the name of the property (used as key in property\
            dict)
            leaf_func: the function to compute the property on the leaves
            inode_func: the function to compute the property on the inner\
            nodes (assuming it is already computed on the leaves)

        Parameters of leaf_func:

        * leaf (PathNode): the current node

        Parameters of inode_func:

        * inode (PathNode): the current node
        * children (list(PathNode)): the list of children of the current node
        """

        self.root.compute_bottom_up_property(property_name, leaf_func, inode_func)

    def compute_top_down_property(
            self,
            property_name:Union[str,PathTreeProperty],
            root_func:Callable[["PathNode"], Any],
            notroot_func:Callable[["PathNode", "PathNode"], Any]
            ) -> None:
        """ Compute a property of top-down type in the subtree of the node.

        This function calls the namesake function of the `PathNode` on the root.
        Check `PathNode.compute_top_down_property()` for more details.

        Parameters:
            property_name: the name of the property (used as key in property\
            dict)
            root_func: the function to compute the property on the root
            notroot_func: the function to compute the property on the not-root\
            nodes (assuming it is already computed on the parent)

        Parameters of root_func:

        * root (PathNode): the current node

        Parameters of notroot_func:

        * node (PathNode): the current node
        * parent (PathNode): the parent of the current node
        """

        self.root.compute_top_down_property(property_name, root_func, notroot_func)

    def compute_individual_property(
            self,
            property_name:Union[str,PathTreeProperty],
            property_func:Callable[["PathNode"], Any]
        ):
        """ Compute an individual property in the tree.

        This function calls the namesake function of the `PathNode` on the root.
        Check `PathNode.compute_individual_property()` for more details.

        Parameters:
            property_name: the name of the property (used as key in property\
            dict)
            property_func: the function to compute the property on the node

        Parameters of property_func:

        * node: the current node (as PathNode)
        """

        self.root.compute_individual_property(property_name, property_func)

    def compute_basic_property(self, property_name:PathTreeProperty) -> None:
        """ Compute a basic property, contained in `PathTreeProperty` enum.

        If the property specified is PathTreeProperty.PRUNED, this function is
        a nop.

        Parameters:
            property_name: the name of the basic property to compute
        """

        if property_name == PathTreeProperty.HEIGHT:
            self.__compute_height()
        elif property_name == PathTreeProperty.DEPTH:
            self.__compute_depth()
        elif property_name == PathTreeProperty.NUM_OF_DIRECTORIES:
            self.__compute_num_of_directories()
        elif property_name == PathTreeProperty.NUM_OF_FILES:
            self.__compute_num_of_files()
        elif property_name == PathTreeProperty.NUM_OF_INODES:
            self.__compute_num_of_inode()
        elif property_name == PathTreeProperty.NUM_OF_LEAVES:
            self.__compute_num_of_leaves()
        elif property_name == PathTreeProperty.NUM_OF_NODES:
            self.__compute_num_of_nodes()
        elif property_name == PathTreeProperty.SIZE:
            self.__compute_size()
        elif property_name == PathTreeProperty.SIMPLE_SIZE:
            self.__compute_simple_size()
        elif property_name == PathTreeProperty.PRUNED:
            pass

    def __compute_height(self):
        """ Compute the height of the nodes in the tree.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.HEIGHT,
            lambda leaf: 0,
            lambda inode, children:
                1 + min(list(int(child.property[PathTreeProperty.HEIGHT])
                    for child in children
                ))
        )

    def __compute_depth(self):
        """ Compute the depth of the nodes in the tree.
        """

        self.compute_top_down_property(
            PathTreeProperty.DEPTH,
            lambda root: 0,
            lambda node, parent: 1 + parent.property[PathTreeProperty.DEPTH]
        )

    def __compute_num_of_directories(self):
        """ Compute the num. of directories in all the subtrees.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.NUM_OF_DIRECTORIES,
            lambda leaf: 1 if leaf.path.is_dir() else 0,
            lambda inode, children:
                (1 if inode.path.is_dir() else 0) +
                sum(list(child.property[PathTreeProperty.NUM_OF_DIRECTORIES]
                    for child in children
                ))
        )

    def __compute_num_of_files(self):
        """ Compute the num. of files in all the subtrees.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.NUM_OF_FILES,
            lambda leaf: 1 if leaf.path.is_file() else 0,
            lambda inode, children:
                sum(list(child.property[PathTreeProperty.NUM_OF_FILES]
                    for child in children
                ))
        )

    def __compute_num_of_inode(self):
        """ Compute the num. of internal nodes in all the subtrees.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.NUM_OF_INODES,
            lambda leaf: 0,
            lambda inode, children:
                1 + sum(list(child.property[PathTreeProperty.NUM_OF_INODES]
                    for child in children
                ))
        )

    def __compute_num_of_leaves(self):
        """ Compute the num. of leaves in all the subtrees.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.NUM_OF_LEAVES,
            lambda leaf: 1,
            lambda inode, children:
                sum(list(child.property[PathTreeProperty.NUM_OF_LEAVES]
                    for child in children
                ))
        )

    def __compute_num_of_nodes(self):
        """ Compute the num. of nodes in all the subtrees.
        """

        self.compute_individual_property(
            PathTreeProperty.NUM_OF_NODES,
            lambda node:
                node.property[PathTreeProperty.NUM_OF_DIRECTORIES] +
                node.property[PathTreeProperty.NUM_OF_FILES]
        )

    def __compute_size(self):
        """ Compute the size (in bytes) of all the subtrees.
        """

        self.compute_bottom_up_property(
            PathTreeProperty.SIZE,
            lambda leaf: leaf.path.stat().st_size if leaf.path.is_file() else 0,
            lambda inode, children:
                sum(list(child.property[PathTreeProperty.SIZE]
                    for child in children
                ))
        )

    def __compute_simple_size(self):
        """ Compute the (simplified) size of all the subtrees.

        The size expressed with a multiple of the byte and with 1-3 significant
        digits. E.g. 23 MB
        """

        def simplified_size(node:PathNode):
            size = node.property[PathTreeProperty.SIZE]
            return \
                f"{size} B" if size < Size.KB.value else (
                f"{size // Size.KB.value} KB" if size < Size.MB.value else (
                f"{size // Size.MB.value} MB" if size < Size.GB.value else (
                f"{size // Size.GB.value} GB" if size < Size.TB.value else
                f"{size // Size.TB.value} TB"
                )))

        self.compute_individual_property(
            PathTreeProperty.SIMPLE_SIZE,
            simplified_size
        )

    def __compute_basic_properties(self) -> None:
        """ Compute the basic properties for the nodes of the tree.

        The basic properties are:

        * height of the nodes
        * depth of the nodes
        * number of directories in the nodes' subtrees
        * number of files in the nodes' subtrees
        * number of nodes in the nodes' subtrees
        * size (in bytes) of the nodes' subtrees
        * simple size (in KB...TB) of the nodes' subtrees
        """

        self.compute_basic_property(PathTreeProperty.HEIGHT)
        self.compute_basic_property(PathTreeProperty.DEPTH)
        self.compute_basic_property(PathTreeProperty.NUM_OF_DIRECTORIES)
        self.compute_basic_property(PathTreeProperty.NUM_OF_FILES)
        self.compute_basic_property(PathTreeProperty.NUM_OF_INODES)
        self.compute_basic_property(PathTreeProperty.NUM_OF_LEAVES)
        self.compute_basic_property(PathTreeProperty.NUM_OF_NODES)
        self.compute_basic_property(PathTreeProperty.SIZE)
        self.compute_basic_property(PathTreeProperty.SIMPLE_SIZE)

    def remove_property(self, property_name:Union[str,PathTreeProperty]) -> tuple[bool, bool]:
        """ Remove a property from all the nodes in the tree.

        If the property is missing from one node, no exception is raised.

        Parameters:
            property_name: the name of the property to remove.

        Returns:
            A tuple containing two booleans:

            * True if the property appeared in all nodes, false otherwise
            * True if the property appeared in at least one node, false\
            otherwise
        """

        in_all_nodes = True
        in_one_node = False
        for node in self:
            status = node.remove_property(property_name)
            in_all_nodes = in_all_nodes and status
            in_one_node = in_one_node or status
        return in_all_nodes, in_one_node

    def physical_pruning(self, keep_condition:Callable[[PathNode],bool]) -> None:
        """ Remove (physically) the all the subtrees where the root does not
        satisfy the keep condition.

        The tree is scanned in breadth-first order. For each node, the keep
        condition is checked and if it is not satisfied all the corresponding
        subtree is physically removed from the tree.

        Note that the properties of the nodes are not recomputed.

        Parameters:
            keep_condition: the boolean function that assess if a node, and its\
            subtree, should be kept or pruned.

        Parameters of keep_condition:

        * node: the node to check.

        Return of keep_condition:

        * True if the node (and subtree) must be kept, false if it must be\
            pruned.
        """

        nodes = [self.root]
        while len(nodes) > 0:
            node = nodes.pop(0)

            idx = 0
            while idx < len(node.children):
                if not keep_condition(node.children[idx]):
                    node.children.pop(idx)
                else:
                    idx += 1

            nodes = nodes + node.children

    def logical_pruning(self, keep_condition:Callable[[PathNode],bool]) -> None:
        """ Remove (logically) the all the subtrees where the root does not
        satisfy the keep condition.

        The logical removal is applied using the property
        `PathTreeProperty.PRUNED`: if true, the node is removed.

        The tree is scanned in breadth-first order. For each node, the keep
        condition is checked and if it is not satisfied all the corresponding
        subtree is logically removed from the tree.

        Parameters:
            keep_condition: the boolean function that assess if a node, and its\
            subtree, should be kept or pruned.

        Parameters of keep_condition:

        * node: the node to check.

        Return of keep_condition:

        * True if the node (and subtree) must be kept, false if it must be\
        pruned.
        """

        self.root.compute_top_down_property(
            PathTreeProperty.PRUNED,
            lambda root: False,
            lambda node, parent:
                parent.property[PathTreeProperty.PRUNED] or
                not keep_condition(node)
        )

    def get_node(self, path:Union[str,Path]) -> Union[PathNode, None]:
        """ Return the PathNode corresponding to the passed Path.

        Parameters:
            path: the path to search in the tree

        Returns:
            The PathNode corresponding to the passed Path if exists, None\
            otherwise.
        """

        if isinstance(path, str):
            path = Path(path)
        parts = path.parts
        node = self.root

        idx = 1
        if node.path == Path("."):
            idx = 0

        found = True
        while found:
            if path == node.path:
                return node
            found = False
            for child in node.children:
                print(f"compare part {idx} between '{parts}' and '{child.path.parts}'")
                if parts[idx] == child.path.parts[idx]:
                    node = child
                    found = True
                    break
            idx += 1
        return None

    def copy(self) -> "PathTree":
        """ Return a deepcopy of the tree and all its nodes.

        Returns:
            A deepcopy of the tree.
        """

        new_root = self.root.copy()
        new_nodes = [new_root]
        while len(new_nodes) > 0:
            new_node = new_nodes.pop(0)
            new_children = [child.copy() for child in new_node.children]
            new_node.children = new_children
            for new_child in new_children:
                new_child.parent = new_node
            new_nodes = new_nodes + new_children
        return PathTree(new_root, skip_properties=True)

    def __str__(self) -> str:
        """ Return a string describing the tree properties.

        Print the property of the root.

        Returns:
            A string describing the tree.
        """

        return f"[PathTree] > Root: ({str(self.root)})"

    def to_csv(
            self,
            csvfile:Union[Path, str],
            properties:Union[list[Union[str, PathTreeProperty]], None]=None,
            node_condition:Callable[[PathNode], bool]=lambda node: True,
            node_limit:int=1000000
        ) -> None:
        """ Export all nodes of the tree to a csv.

        The export includes the name of the path and a list of properties. Due
        to the high number of nodes a directory tree can have, by default, the
        export is limited to the first 1 million nodes.

        Parameters:
            csvfile: the name of the csv for the export.
            properties: the list of properties to include in the export. If\
            None all parameters are included.
            node_condition: the condition a node must meet to be exported (by\
            default all nodes are exported).
            node_limit: the max number of nodes that can be exported. If <= 0,\
            no limitation is applied.
        """

        if isinstance(csvfile, str):
            csvfile = Path(csvfile)
        if properties is None:
            properties = list(self.root.property.keys())

        # Header
        str_properties = []
        for prop in properties:
            if isinstance(prop, PathTreeProperty):
                str_properties.append(prop.value)
        header = ";".join(["path"] + str_properties)

        # Data
        lines = []
        for node in self.breadth_first_iter():
            if 0 < node_limit <= len(lines):
                break
            if not node_condition(node):
                continue
            line = [node.path.as_posix()]
            for curr_property in properties:
                line.append(str(node.property[curr_property]))
            lines.append(";".join(line))

        with open(csvfile, "w", encoding="utf8") as f:
            f.write(header + "\n")
            for line in lines:
                f.write(line + "\n")

    def to_excel(
            self,
            excelfile:Union[Path, str],
            properties:Union[list[Union[str, PathTreeProperty]], None]=None,
            node_condition:Callable[[PathNode], bool]=lambda node: True,
            node_limit:int=1000000
        ) -> None:
        """ Export all nodes of the tree to Excel.

        The export includes the name of the path and a list of properties. Due
        to the high number of nodes a directory tree can have, by default, the
        export is limited to the first 1 million nodes.

        Parameters:
            excelfile: the name of the Excel for the export.
            properties: the list of properties to include in the export. If\
            None all parameters are included.
            node_condition: the condition a node must meet to be exported (by\
            default all nodes are exported).
            node_limit: the max number of nodes that can be exported. If <= 0,\
            no limitation is applied.
        """

        if isinstance(excelfile, str):
            excelfile = Path(excelfile)

        if properties is None:
            properties = list(self.root.property.keys())

        workbook = xlsxwriter.Workbook(excelfile)
        sheet = workbook.add_worksheet(f"Export - {self.root.path}")

        # Header
        header_format = workbook.add_format({
            'bold': True,
            'font_color': 'white',
            'bg_color': "#002060"
        })
        sheet.write(0, 0, "Path", header_format)
        for col, prop in enumerate(properties, 1):
            if isinstance(prop, PathTreeProperty):
                sheet.write(0, col, prop.value, header_format)
            else:
                sheet.write(0, col, prop, header_format)

        # Data
        row = 1
        for node in self.breadth_first_iter():
            if 0 < node_limit < row:
                break
            if not node_condition(node):
                continue

            sheet.write(row, 0, str(node.path))
            for col, prop in enumerate(properties, 1):
                if prop in node.property:
                    sheet.write(row, col, node.property[prop])
            row += 1

        sheet.autofit()
        workbook.close()


#!/bin/env python
# This is a script to use descriptive statistics calculated previously to produce presentation-quality graphics
# and save them as PNG files for use in a powerpoint presentation
#   written by Justin Meyer
#   last edited 2022-05-04
#

# import required modules
import pandas as pd
import matplotlib.pyplot as plt

import os
os.chdir('c:/Users/justi/OneDrive/Documents/GitHub/final-project-meyer443')

def ReadData( fileName ):
    """This function takes a filename as input, and returns a dataframe with
    raw data read from that file in a Pandas DataFrame."""

    # open and read the file
    DataDF = pd.read_csv( fileName, sep=',' )
    return DataDF

def Plot_Temp(DataDF, outFileName):
    plt.scatter(DataDF["Date"], DataDF["temp"], marker=".") # scatter date vs temp
    plt.xlabel('Date')
    plt.xticks([]) # can't read x axis with overlapping dates
    plt.ylabel('Average Daily Temperature (C)')
    plt.savefig('Daily_Temp.png', dpi=96) # save the plot as PNG with 96 dpi
    plt.close()

# the following condition checks whether we are running as a script, in which
# case run the test code, otherwise functions are being imported so do not.
# put the main routines from your code after this conditional check.

if __name__ == '__main__':

    fileName = "Daily_Temperature_Averages.csv"

    DataDF = ReadData( fileName )

# Figure generation
    outTempName = "Daily_Temp.png"
    Plot_Temp( DataDF, outTempName )