from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import sys
import os

import time

import numpy as np

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

from ..smplifyx.utils_mics.mesh_viewer import MeshViewer
from ..smplifyx import utils_mics
from ..smplifyx.utils_mics import misc_utils
from ..smplifyx import fitting as single_view_fitting
from ..constants import SKELETON_IDX, HAND_IDX
# from .fitting_video_orientation_temporal_loss import SMPLifyBodyOrientLoss, MultiViewSMPLifyBodyOrientLoss
from loguru import logger

from .body_pose_utils import BODY_POSE_TO_OP, FEET_IN_SMPL, FEET_IN_OP
# from .fitting_video_persubject_orientation_temporal import *
# ! SMPLifyX: init in OpenGL CS.

class SMPLifyLoss3D(single_view_fitting.SMPLifyLoss):

    def __init__(self, search_tree=None,
                 pen_distance=None, tri_filtering_module=None,
                 rho=100,
                 body_pose_prior=None,
                 shape_prior=None,
                 angle_prior=None,
                 dtype=torch.float32,
                 data_weight=1.0,
                 body_pose_weight=0.0,
                 shape_weight=0.0,
                 bending_prior_weight=0.0,
                 hand_prior_weight=0.0,
                 expr_prior_weight=0.0, jaw_prior_weight=0.0,
                 reduction='sum',
                 **kwargs):

        super(SMPLifyLoss3D, self).__init__(search_tree=search_tree,
                 pen_distance=pen_distance, tri_filtering_module=tri_filtering_module,
                 rho=rho,
                 body_pose_prior=body_pose_prior,
                 shape_prior=shape_prior,
                 angle_prior=angle_prior,
                 dtype=torch.float32,
                 data_weight=data_weight,
                 body_pose_weight=body_pose_weight,
                 shape_weight=shape_weight,
                 bending_prior_weight=bending_prior_weight,
                 reduction='sum',
                 **kwargs)
        # import pdb;pdb.set_trace()
        self.ground_plane_support = kwargs['ground_plane_support']
        self.gp_support_loss_weight = kwargs['gp_support_loss_weight_init']

    def forward(self, body_model_output, gt_joints,
                body_model_faces,
                use_vposer=False, pose_embedding=None,
                scene_model=None,
                **kwargs):

        batch_size = gt_joints.shape[0]
        # import pdb;pdb.set_trace()
        # Calculate the distance of the projected joints from
        # the ground truth 2D detections

        #### OLD used in CVPR submission.
        # joint_diff = (gt_joints - body_model_output.joints[:,:15,:]) ** 2 / batch_size
        # joint_loss = (torch.sum(joint_diff) * self.data_weight ** 2) / batch_size
        #### end of OLD used in CVPR submission.

        logger.info('run 3D joint loss.')
        # delete face contour in original name
        joint_len = min(gt_joints.shape[1], body_model_output.joints.shape[1])

        print('joints shape', gt_joints.shape[-1])
        # TODO: check the idx.
        if gt_joints.shape[-1] == 3:
            joint_diff = (gt_joints[:, :joint_len, :] - body_model_output.joints[:,:joint_len,:]) ** 2
        elif gt_joints.shape[-1] == 4:
            joint_diff = (gt_joints[:, :joint_len, :-1] - body_model_output.joints[:,:joint_len,:]) ** 2
            joint_diff *= gt_joints[:, :joint_len, -1:]

        joint_loss = (torch.sum(joint_diff) * self.data_weight ** 2)

        # visualizaion: use updated translated 3D joint in Perspective Camera.
        # Render Joint is still not accurate with the original joints.
        if False:
            import pdb;pdb.set_trace()
            cameras = kwargs['cameras']

            import cv2
            img_path  = '//ps/scratch/hyi/HCI_dataset/20210109_capture/C0034/hci_test/000379.jpg'
            cv2_img = cv2.imread(img_path)
            cam = cameras[0]
            proj_kpts = cam(gt_joints)
            ori_proj = proj_kpts.squeeze(0).detach().cpu().numpy().T
            for j in range(ori_proj.shape[1]):
                print(ori_proj[0, j], ori_proj[1, j])
                cv2_img = cv2.circle(cv2_img, (int(ori_proj[0, j]), int(ori_proj[1, j])), \
                            radius=5, color=(255, 0, 0), thickness=4)

            proj_kpts = cam(body_model_output.joints[:,:15,:])
            ori_proj = proj_kpts.squeeze(0).detach().cpu().numpy().T
            for j in range(ori_proj.shape[1]):
                print(ori_proj[0, j], ori_proj[1, j])
                cv2_img = cv2.circle(cv2_img, (int(ori_proj[0, j]), int(ori_proj[1, j])), \
                            radius=5, color=(0, 0, 255), thickness=4)

            # cv2.imshow("2d joint", cv2_img)
            cv2.imwrite("render.png", cv2_img)
            cv2.waitKey(5)
            import pdb;pdb.set_trace()

        # Calculate the loss from the Pose prior
        if use_vposer:
            pprior_loss = (pose_embedding.pow(2).sum() *
                           self.body_pose_weight ** 2)
        else:
            pprior_loss = torch.sum(self.body_pose_prior(
                body_model_output.body_pose,
                body_model_output.betas)) * self.body_pose_weight ** 2
        pprior_loss = pprior_loss / batch_size

        shape_loss = torch.sum(self.shape_prior(
            body_model_output.betas)) * self.shape_weight ** 2 / batch_size

        # Calculate the prior over the joint rotations. This a heuristic used
        # to prevent extreme rotation of the elbows and knees
        body_pose = body_model_output.full_pose[:, 3:66]
        angle_prior_loss = torch.sum(
            self.angle_prior(body_pose)) * self.bending_prior_weight  / batch_size

        total_loss = (joint_loss + pprior_loss + shape_loss + angle_prior_loss)

        message = f'Init joint|weight: {joint_loss}|{self.data_weight.item():.3f},  \
            shape_loss: {shape_loss}|{self.shape_weight}, \
            pprior_loss: {pprior_loss}|{self.body_pose_weight}, \
            angle_prior: {angle_prior_loss}|{self.bending_prior_weight:.3f}, \
            total_loss: {total_loss}'

        # Delete gp plane
        # if self.ground_plane_support:
        #     ## ground plane support loss: human should above the ground plane.
        #     assert scene_model is not None
        #     gt_ground_plane = scene_model.ground_plane
        #     # batch average
        #     # gp_support_loss = F.smooth_l1_loss(torch.max(body_model_output.vertices[:, :, 1]), gt_ground_plane)
        #     # gp_support_loss = gp_support_loss * self.gp_support_loss_weight
        #     gp_support_loss = F.relu(torch.max(body_model_output.vertices[:, :, 1], -1)[0] - gt_ground_plane).mean()
        #     total_loss = total_loss + gp_support_loss * self.gp_support_loss_weight
        #     message += f'gp: {gp_support_loss}'

        logger.info(message)

        tb_debug = kwargs['tb_debug']
        if tb_debug:
            debug_loss_dict = {
                'loss_joint': joint_loss,
                'loss_pprior': pprior_loss,
                'loss_shape': shape_loss,
                'loss_angle_prior': angle_prior_loss,
                'loss_total_SMPLifyLoss3D': total_loss,
            }
            # if self.ground_plane_support:
            #     debug_loss_dict['loss_gp_support_body'] = gp_support_loss
            return total_loss, debug_loss_dict
        else:
            return total_loss

####################
### body skeletons add hands.
####################
class SMPLifyLoss3D_withHands(single_view_fitting.SMPLifyLoss):
    def __init__(self, search_tree=None,
                 pen_distance=None, tri_filtering_module=None,
                 rho=100,
                 body_pose_prior=None,
                 shape_prior=None,
                 angle_prior=None,
                 dtype=torch.float32,
                 data_weight=1.0,
                 body_pose_weight=0.0,
                 shape_weight=0.0,
                 bending_prior_weight=0.0,
                 hand_prior_weight=0.0,
                 reduction='sum',
                 hand_joints_weights=1.0,
                 expr_prior_weight=0.0,
                 jaw_prior_weight=0.0,
                 hand_idx=None,
                 **kwargs):

        super(SMPLifyLoss3D_withHands, self).__init__(search_tree=search_tree,
                 pen_distance=pen_distance, tri_filtering_module=tri_filtering_module,
                 rho=rho,
                 body_pose_prior=body_pose_prior,
                 shape_prior=shape_prior,
                 angle_prior=angle_prior,
                 dtype=torch.float32,
                 data_weight=data_weight,
                 body_pose_weight=body_pose_weight,
                 shape_weight=shape_weight,
                 bending_prior_weight=bending_prior_weight,
                 hand_prior_weight=hand_prior_weight,
                 expr_prior_weight=expr_prior_weight,
                 jaw_prior_weight=jaw_prior_weight,
                 reduction='sum',
                 **kwargs)
        self.hand_joints_weights = hand_joints_weights

        # self.hand_idx_map=hand_idx
        # import pdb;pdb.set_trace()
    def forward(self, body_model_output, gt_joints,
                body_model_faces,
                use_vposer=False, pose_embedding=None,
                scene_model=None,
                **kwargs):

        batch_size = gt_joints.shape[0]
        # import pdb;pdb.set_trace()
        # Calculate the distance of the projected joints from
        # the ground truth 2D detections

        #### OLD used in CVPR submission.
        # joint_diff = (gt_joints - body_model_output.joints[:,:15,:]) ** 2 / batch_size
        # joint_loss = (torch.sum(joint_diff) * self.data_weight ** 2) / batch_size
        #### end of OLD used in CVPR submission.

        logger.info('run 3D joint loss.')
        # delete face contour in original name
        joint_len = min(gt_joints.shape[1], body_model_output.joints.shape[1])

        print('joints shape', gt_joints.shape[-1])
        # TODO: check the idx.
        if gt_joints.shape[-1] == 3:
            joint_diff = (gt_joints[:, :joint_len, :] - body_model_output.joints[:,:joint_len,:]) ** 2
        elif gt_joints.shape[-1] == 4:
            joint_diff = (gt_joints[:, :joint_len, :-1] - body_model_output.joints[:,:joint_len,:]) ** 2
            joint_diff *= gt_joints[:, :joint_len, -1:]

        # TODO: check the data_weight.
        # import pdb;pdb.set_trace()
        ske_loss = (torch.sum(joint_diff[:, SKELETON_IDX]) * self.data_weight ** 2)
        hand_loss = (torch.sum(joint_diff[:, HAND_IDX]) * self.hand_joints_weights ** 2)
        joint_loss = ske_loss + hand_loss
        # Calculate the loss from the Pose prior
        if use_vposer:
            pprior_loss = (pose_embedding.pow(2).sum() *
                           self.body_pose_weight ** 2)
        else:
            pprior_loss = torch.sum(self.body_pose_prior(
                body_model_output.body_pose,
                body_model_output.betas)) * self.body_pose_weight ** 2

        pprior_loss = pprior_loss / batch_size

        shape_loss = torch.sum(self.shape_prior(
            body_model_output.betas)) * self.shape_weight ** 2 / batch_size

        # Calculate the prior over the joint rotations. This a heuristic used
        # to prevent extreme rotation of the elbows and knees
        body_pose = body_model_output.full_pose[:, 3:66]
        angle_prior_loss = torch.sum(
            self.angle_prior(body_pose)) * self.bending_prior_weight  / batch_size

        # Apply the prior on the pose space of the hand
        # import pdb;pdb.set_trace()
        left_hand_prior_loss, right_hand_prior_loss = 0.0, 0.0
        if self.use_hands and self.left_hand_prior is not None:
            left_hand_prior_loss = torch.sum(
                self.left_hand_prior(
                    body_model_output.left_hand_pose)) * \
                self.hand_prior_weight ** 2 / batch_size

        if self.use_hands and self.right_hand_prior is not None:
            right_hand_prior_loss = torch.sum(
                self.right_hand_prior(
                    body_model_output.right_hand_pose)) * \
                self.hand_prior_weight ** 2 / batch_size

        total_loss = (joint_loss + pprior_loss + shape_loss + angle_prior_loss \
                + right_hand_prior_loss + left_hand_prior_loss)

        message = f'skeleton: {ske_loss}, hand: {hand_loss}, \
            joint|weight: {joint_loss}|{self.data_weight.item():.3f},  \
            shape_loss: {shape_loss}|{self.shape_weight}, \
            pprior_loss: {pprior_loss}|{self.body_pose_weight}, \
            angle_prior: {angle_prior_loss}|{self.bending_prior_weight:.3f}, \
            left_hand_prior: {left_hand_prior_loss}|{self.hand_prior_weight:.3f}, \
            total_loss: {total_loss}'

        # Delete gp plane
        # if self.ground_plane_support:
        #     ## ground plane support loss: human should above the ground plane.
        #     assert scene_model is not None
        #     gt_ground_plane = scene_model.ground_plane
        #     # batch average
        #     # gp_support_loss = F.smooth_l1_loss(torch.max(body_model_output.vertices[:, :, 1]), gt_ground_plane)
        #     # gp_support_loss = gp_support_loss * self.gp_support_loss_weight
        #     gp_support_loss = F.relu(torch.max(body_model_output.vertices[:, :, 1], -1)[0] - gt_ground_plane).mean()
        #     total_loss = total_loss + gp_support_loss * self.gp_support_loss_weight
        #     message += f'gp: {gp_support_loss}'

        logger.info(message)

        tb_debug = kwargs['tb_debug']
        if tb_debug:
            debug_loss_dict = {
                'loss_joint': joint_loss,
                'loss_pprior': pprior_loss,
                'loss_shape': shape_loss,
                'loss_angle_prior': angle_prior_loss,
                'loss_total_SMPLifyLoss3D': total_loss,
                'loss_left_hand_prior': left_hand_prior_loss,
                'loss_right_hand_prior': right_hand_prior_loss,
            }
            # if self.ground_plane_support:
            #     debug_loss_dict['loss_gp_support_body'] = gp_support_loss
            return total_loss, debug_loss_dict
        else:
            return total_loss

################################
### based on 2D observation.
################################
# TODO: add online POSA.
class MultiViewSMPLifyLoss(single_view_fitting.SMPLifyLoss):
    def __init__(self, **kwargs):

        super(MultiViewSMPLifyLoss, self).__init__(**kwargs)
        # import pdb;pdb.set_trace()
        ## ground_plane_support
        self.ground_plane_support = kwargs['ground_plane_support']
        if self.ground_plane_support:
            self.register_buffer('gp_support_loss_weight',
                                    torch.tensor(kwargs['gp_support_loss_weight'], dtype=kwargs['dtype']))

        ## ground_contact_support
        self.ground_contact_support = kwargs['ground_contact_support']
        if self.ground_contact_support:
            self.register_buffer('gp_contact_loss_weight',
                                    torch.tensor(kwargs['gp_contact_loss_weight'], dtype=kwargs['dtype']))
            self.ground_contact_vertices_ids = kwargs['ground_contact_vertices_ids']
            self.rho_contact = kwargs['rho_contact']
            logger.info(f'rho_contact in gp_contact_robustifier: {kwargs["rho_contact"]}')
            self.gp_contact_robustifier = misc_utils.GMoF_unscaled(rho=self.rho_contact)

        ## human & scene object inter-penetration
        self.sdf_penetration = kwargs['sdf_penetration']
        if self.sdf_penetration:
            self.register_buffer('sdf_penetration_loss_weight',
                                    torch.tensor(kwargs['sdf_penetration_loss_weight'], dtype=kwargs['dtype']))

        ## human & scene object contact loss
        self.contact = kwargs['contact']
        if self.contact or True:
            self.register_buffer('contact_loss_weight',
                                    torch.tensor(kwargs['contact_loss_weight'], dtype=kwargs['dtype']))
            self.contact_verts_ids = kwargs['contact_verts_ids']
            self.rho_contact = kwargs['rho_contact']
            self.contact_angle = kwargs['contact_angle']
            self.ftov = kwargs['ftov']
            self.contact_robustifier = misc_utils.GMoF(rho=self.rho_contact) # rho is the distance. 5cm
        # else:
        #     self.contact_verts_ids = None
        #     self.rho_contact = None
        #     self.contact_angle = None
        #     self.ftov = None
        # self.contact_robustifier = misc_utils.GMoF_unscaled(rho=self.rho_contact) #used in prox

        ## depth ordinal loss
        self.ordinal_depth = kwargs['ordinal_depth']
        if self.ordinal_depth:
            self.register_buffer('ordinal_depth_loss_weight', torch.tensor(kwargs['ordinal_depth_loss_weight'], dtype=kwargs['dtype']))

        self.scene = kwargs['scene']
        if self.scene:
            self.register_buffer('scene_loss_weight',
                                    torch.tensor(kwargs['scene_loss_weight'], dtype=kwargs['dtype']))

        self.loss_use_sum = kwargs['loss_use_sum']
        self.video_smooth = kwargs['video_smooth']
        self.constant_velocity = kwargs['constant_velocity']
        # learn from AMASS dataset
        self.motion_smooth_prior_flag = kwargs['motion_smooth_prior']

        if self.video_smooth:
            self.register_buffer('smooth_2d_weight',
                                    torch.tensor(kwargs['smooth_2d_weight'], dtype=kwargs['dtype']))
            self.register_buffer('smooth_3d_weight',
                                    torch.tensor(kwargs['smooth_3d_weight'], dtype=kwargs['dtype']))

            self.robustifier_3DJoint = misc_utils.GMoF(rho=0.1)

            if self.motion_smooth_prior_flag:
                self.register_buffer('motion_prior_weight',
                                    torch.tensor(kwargs['motion_prior_weight'], dtype=kwargs['dtype']))
                sys.path.append('/is/cluster/hyi/workspace/HCI/amass_infiller')
                from motion_smooth_prior import MotionSmoothPrior
                self.motion_smooth_prior = MotionSmoothPrior()

        ## pare pose loss weight
        self.pare_pose_prior = kwargs['pare_pose_prior']
        if self.pare_pose_prior:
            self.register_buffer('pare_pose_weight',
                                torch.tensor(kwargs['pare_pose_weight'], dtype=kwargs['dtype']))

        ## for ordinal depth loss
        self.use_human_depth = kwargs['use_human_depth']

        self.depth_robustifier = None # ! without using it, sometimes, it will fail.
        # self.depth_robustifier = kwargs['depth_robustifier'] # default is None, set 10

        # gp contact loss
        self.rho_gp_flag = None
        self.l1_gp_flag = None

        # posa model for online pose

        # posa pre-calculated contact labels
        self.posa_flag=kwargs['posa_flag']
        self.online_pose=kwargs['online_pose']
        self.posa_body_contact_labels = kwargs['posa_body_contact_labels']

        # codebase selection
        self.use_scene_loss = kwargs['use_scene_loss']
        # import pdb;pdb.set_trace()
        self.save_dir = kwargs['output_folder']

        # TODO: add motion smooth loss term.

    def forward(self, body_model_output, cameras, gt_joints, joints_conf,
                body_model_faces, joint_weights,
                use_vposer=False, pose_embedding=None,
                gt_contact_value=None,
                scene_model=None,
                perframe_mask=None,
                pare_body_pose=None, pare_body_flag=None,
                **kwargs):

        ################################
        # joints_conf: smpl format;
        # gt_joints: openpose 2D joints;
        ################################
        batch_size = gt_joints.shape[0]

        # TODO: the shape is not correct !!!
        gt_joints = gt_joints.transpose(0, 1)
        joints_conf = joints_conf.unsqueeze(0)

        joint_loss = 0

        # for motion smooth loss term.
        smooth_j2d_loss = 0
        smooth_j3d_loss = 0
        motion_prior_loss = 0
        for vi, cam in enumerate(cameras):
            projected_joints = cam(body_model_output.joints)
            # Calculate the weights for each joints
            # batch * n_joints_1

            if False:
                import cv2
                img_path  = '//ps/scratch/hyi/HCI_dataset/20210109_capture/C0034/hci_test/000379.jpg'
                cv2_img = cv2.imread(img_path)
                cam = cameras[0]
                ori_proj = gt_joints[vi, :,:15,:].squeeze(0).detach().cpu().numpy().T
                for j in range(ori_proj.shape[1]):
                    print(ori_proj[0, j], ori_proj[1, j])
                    cv2_img = cv2.circle(cv2_img, (int(ori_proj[0, j]), int(ori_proj[1, j])), \
                                radius=5, color=(255, 0, 0), thickness=4)

                proj_kpts = cam(body_model_output.joints[:,:15,:])
                ori_proj = proj_kpts.squeeze(0).detach().cpu().numpy().T
                for j in range(ori_proj.shape[1]):
                    print(ori_proj[0, j], ori_proj[1, j])
                    cv2_img = cv2.circle(cv2_img, (int(ori_proj[0, j]), int(ori_proj[1, j])), \
                                radius=5, color=(0, 0, 255), thickness=4)

                # cv2.imshow("2d joint", cv2_img)
                cv2.imwrite("mvs_render.png", cv2_img)
                cv2.waitKey(5)
                import pdb;pdb.set_trace()

            weights = (joint_weights * joints_conf[vi]
                       if self.use_joints_conf else
                       joint_weights).unsqueeze(dim=-1)

            # Calculate the distance of the projected joints from
            # the ground truth 2D detections
            # TODO: the shape is not correct !!!
            joint_diff = self.robustifier(gt_joints[vi, :, :, :] - projected_joints)
            joint_loss = joint_loss + (torch.sum(weights ** 2 * joint_diff) *
                        self.data_weight ** 2) / batch_size

            if self.video_smooth:
                if self.motion_smooth_prior_flag:
                    # TODO:
                    pass
                    model_verts = body_model_output.vertices
                    model_joints = body_model_output.joints
                    # * factor out missing frame effects.
                    # import pdb;pdb.set_trace()
                    motion_useful_conf = (joints_conf[vi, 1:, :].sum(-1) > 0) & (joints_conf[vi, 0:-1, :].sum(-1) > 0)
                    motion_prior_loss = self.motion_smooth_prior(model_verts, model_joints, conf=motion_useful_conf) # per-single image loss.
                    motion_prior_loss = self.motion_prior_weight * motion_prior_loss

                if self.constant_velocity:
                    # import pdb;pdb.set_trace()
                    if True:
                        ## smooth 2D joint loss
                        joint_conf_diff = joints_conf[vi, 1:-1, :]
                        joints_2d_diff = self.robustifier(projected_joints[2:]+projected_joints[:-2] - 2 * projected_joints[1:-1])
                        weights_diff = (joint_weights * joint_conf_diff
                                if self.use_joints_conf else
                                joint_conf_diff)
                        smooth_j2d_loss = (weights_diff ** 2) * joints_2d_diff.sum(dim=-1)
                        smooth_j2d_loss = torch.cat(
                                [torch.zeros(1, smooth_j2d_loss.shape[1], device=joints_conf.device), smooth_j2d_loss]
                            ).sum(dim=-1)
                        smooth_j2d_loss = (self.smooth_2d_weight ** 2) * smooth_j2d_loss / batch_size
                        smooth_j2d_loss = smooth_j2d_loss.sum()

                        ## smooth 3D joint loss
                        model_joints = body_model_output.joints
                        joints_3d_diff = self.robustifier_3DJoint(model_joints[2:]+model_joints[:-2] - 2 * model_joints[1:-1])
                        smooth_j3d_loss = (weights_diff ** 2) * joints_3d_diff.sum(dim=-1)
                        smooth_j3d_loss = torch.cat(
                                [torch.zeros(1, smooth_j3d_loss.shape[1], device=joints_conf.device), smooth_j3d_loss]
                            ).sum(dim=-1)
                        smooth_j3d_loss = (self.smooth_3d_weight ** 2) * smooth_j3d_loss / batch_size
                        smooth_j3d_loss = smooth_j3d_loss.sum()
                    else: # static motion prior
                        ## smooth 2D joint loss
                        joint_conf_diff = joints_conf[vi, 1:, :]
                        joints_2d_diff = projected_joints[1:] - projected_joints[:-1]
                        weights_diff = (joint_weights * joint_conf_diff
                                if self.use_joints_conf else
                                joint_conf_diff)
                        smooth_j2d_loss = (weights_diff ** 2) * joints_2d_diff.abs().sum(dim=-1)
                        smooth_j2d_loss = torch.cat(
                                [torch.zeros(1, smooth_j2d_loss.shape[1], device=joints_conf.device), smooth_j2d_loss]
                            ).sum(dim=-1)
                        smooth_j2d_loss = (self.smooth_2d_weight ** 2) * smooth_j2d_loss / batch_size
                        smooth_j2d_loss = smooth_j2d_loss.sum()

                        ## smooth 3D joint loss
                        model_joints = body_model_output.joints
                        joints_3d_diff = model_joints[1:] - model_joints[:-1]
                        smooth_j3d_loss = (weights_diff ** 2) * joints_3d_diff.abs().sum(dim=-1)
                        smooth_j3d_loss = torch.cat(
                                [torch.zeros(1, smooth_j3d_loss.shape[1], device=joints_conf.device), smooth_j3d_loss]
                            ).sum(dim=-1)
                        smooth_j3d_loss = (self.smooth_3d_weight ** 2) * smooth_j3d_loss / batch_size
                        smooth_j3d_loss = smooth_j3d_loss.sum()

        pare_pose_prior_loss = 0.0
        if self.pare_pose_prior and self.pare_pose_weight > 0:
            # import pdb;pdb.set_trace()
            assert pare_body_pose.shape[0] == pare_body_flag.shape[0]
            # pose embedding to body pose (angles)
            tmp_body_pose = body_model_output.body_pose
            use_ful = [map_flag != -1 for map_flag in BODY_POSE_TO_OP]
            use_op = [map_flag for map_flag in BODY_POSE_TO_OP if map_flag != -1]
            op_conf_body = torch.ones((batch_size, 21)).type_as(joints_conf) # warining: inverse op conf
            op_conf_body[:, use_ful] = joints_conf[0][:, use_op]

            op_conf_body = op_conf_body.unsqueeze(-1).repeat(1, 1, 3).reshape(batch_size, -1)
            # TODO: modify pare_body_pose detach
            # import pdb;pdb.set_trace()
            # pare_pose_penalty = (1- op_conf_body) * torch.abs(body_model_output.body_pose - pare_body_pose.detach()).pow(2).sqrt()#.reshape(-1, 21, 3).reshape(batch_size, -1)
            pare_pose_penalty = (1 - op_conf_body) * torch.abs(body_model_output.body_pose - pare_body_pose.detach())
            pare_pose_prior_loss = torch.sum(pare_body_flag * pare_pose_penalty, 1)
            pare_pose_prior_loss = (pare_pose_prior_loss * self.pare_pose_weight ** 2).mean()

        # Calculate the loss from the Pose prior
        if use_vposer:
            pprior_loss = (pose_embedding.pow(2).sum() *
                           self.body_pose_weight ** 2)
        else:
            pprior_loss = torch.sum(self.body_pose_prior(
                body_model_output.body_pose,
                body_model_output.betas)) * self.body_pose_weight ** 2
        pprior_loss = pprior_loss / batch_size

        shape_loss = torch.sum(self.shape_prior(
            body_model_output.betas)) * self.shape_weight ** 2 / batch_size
        # Calculate the prior over the joint rotations. This a heuristic used
        # to prevent extreme rotation of the elbows and knees
        body_pose = body_model_output.full_pose[:, 3:66]
        angle_prior_loss = torch.sum(
            self.angle_prior(body_pose)) * self.bending_prior_weight / batch_size

        # Apply the prior on the pose space of the hand
        left_hand_prior_loss, right_hand_prior_loss = 0.0, 0.0
        if self.use_hands and self.left_hand_prior is not None:
            left_hand_prior_loss = torch.sum(
                self.left_hand_prior(
                    body_model_output.left_hand_pose)) * \
                self.hand_prior_weight ** 2 / batch_size

        if self.use_hands and self.right_hand_prior is not None:
            right_hand_prior_loss = torch.sum(
                self.right_hand_prior(
                    body_model_output.right_hand_pose)) * \
                self.hand_prior_weight ** 2 / batch_size

        expression_loss = 0.0
        jaw_prior_loss = 0.0
        if self.use_face:
            expression_loss = torch.sum(self.expr_prior(
                body_model_output.expression)) * \
                self.expr_prior_weight ** 2 / batch_size

            if hasattr(self, 'jaw_prior'):
                jaw_prior_loss = torch.sum(
                    self.jaw_prior(
                        body_model_output.jaw_pose.mul(
                            self.jaw_prior_weight))) / batch_size

        pen_loss = 0.0
        # Calculate the loss due to interpenetration
        if (self.interpenetration and self.coll_loss_weight.item() > 0):

            batch_size = gt_joints.shape[1]
            triangles = torch.index_select(
                body_model_output.vertices, 1,
                body_model_faces).view(batch_size, -1, 3, 3)

            with torch.no_grad():
                collision_idxs = self.search_tree(triangles)

            # Remove unwanted collisions
            if self.tri_filtering_module is not None:
                collision_idxs = self.tri_filtering_module(collision_idxs)

            if collision_idxs.ge(0).sum().item() > 0:
                pen_loss = torch.sum(
                    self.coll_loss_weight *
                    self.pen_distance(triangles, collision_idxs)) / batch_size

        ############## ! warning: calculate in world coordinates system.
        body_vertices_world = body_model_output.vertices

        # gp_support_loss = 0.0
        # if self.ground_plane_support and self.gp_support_loss_weight > 0 and scene_model is not None:
        #     gt_ground_plane = scene_model.ground_plane
        #     ## ground plane support loss
        #     assert gt_ground_plane is not None
        #     gp_support_loss = F.smooth_l1_loss(torch.max(body_vertices_world[:, :, 1]), gt_ground_plane)
        #     gp_support_loss = gp_support_loss * self.gp_support_loss_weight

        # gp_contact_loss = 0.0
        # if self.ground_contact_support and self.gp_contact_loss_weight > 0 and scene_model is not None :
        #     # import pdb;pdb.set_trace()
        #     gt_ground_plane = scene_model.ground_plane
        #     assert gt_contact_value is not None
        #     # the distance from heel and toe to the mesh surface should be
        #     # consist with the distance from joint to the ground plane

        #     # gt_contact_value[gt_contact_value==0] = 0 #-0.2
        #     # gp_contact_loss = torch.abs(body_vertices_world[:, self.ground_contact_vertices_ids, 1] - gt_ground_plane)
        #     # gp_contact_loss = (gt_contact_value * gp_contact_loss.mean(-1) * self.gp_contact_loss_weight).sum() / (gt_contact_value.sum().to(torch.float32) + 1e-9)

        #     gp_contact_loss = F.smooth_l1_loss(body_vertices_world[:, self.ground_contact_vertices_ids, 1].mean(-1), gt_ground_plane, reduction='none')

        #     gp_contact_loss = self.contact_robustifier(gp_contact_loss)
        #     gp_contact_loss = self.gp_contact_loss_weight * gp_contact_loss.mean()

        # if self.ground_contact_support and self.gp_contact_loss_weight > 0 and scene_model is not None :
        #     # import pdb;pdb.set_trace()
        #     gt_ground_plane = scene_model.ground_plane
        #     assert gt_contact_value is not None
        #     # the distance from heel and toe to the mesh surface should be
        #     # consist with the distance from joint to the ground plane

        #     gt_contact_value[gt_contact_value==0] = 0 #-0.2
        #     gp_contact_loss = torch.abs(body_vertices_world[:, self.ground_contact_vertices_ids, 1] - gt_ground_plane)
        #     gp_contact_loss = (gt_contact_value * gp_contact_loss.mean(-1) * self.gp_contact_loss_weight).sum() / (gt_contact_value.sum() + 1e-9)
        gp_support_loss = 0.0
        gp_contact_loss = 0.0

        # import pdb;pdb.set_trace()
        if True: # old version for single image contact loss.
            # import pdb;pdb.set_trace()
            if self.ground_plane_support:
                gt_ground_plane = scene_model.ground_plane
                gp_support_loss = F.relu(torch.max(body_vertices_world[:, :, 1], -1)[0] - gt_ground_plane).mean()
                gp_support_loss = gp_support_loss * self.gp_support_loss_weight

            if self.ground_contact_support and self.gp_contact_loss_weight > 0:
                # contact ids from POSA <-> ground plane
                # TODO: contact_loss(contacted feet <-> ground plane)
                gt_ground_plane = scene_model.ground_plane
                dist_gp = self.gp_contact_robustifier(torch.abs(body_vertices_world[:, :, 1] - gt_ground_plane))
                tmp_feet_label = torch.zeros(self.posa_body_contact_labels.shape).type_as(self.posa_body_contact_labels)
                tmp_feet_label[:, self.ground_contact_vertices_ids.reshape(-1)] = 1
                body_vertices_contact_label_feet = self.posa_body_contact_labels & tmp_feet_label
                # import pdb;pdb.set_trace()
                if body_vertices_contact_label_feet.sum() > 0:
                    gp_contact_loss = dist_gp[body_vertices_contact_label_feet.squeeze(-1)].mean() * self.gp_contact_loss_weight
                else:
                    gp_contact_loss = torch.zeros(1).cuda()

        elif False:
            if self.ground_contact_support and self.gp_contact_loss_weight > 0:
                # import pdb;pdb.set_trace()
                gt_ground_plane = scene_model.ground_plane
                feet_vertices = body_vertices_world[:, self.ground_contact_vertices_ids, 1] # B, 2, N
                # reliable feet: op conf > thre
                # import pdb;pdb.set_trace()
                # joints_conf[vi][:, FEET_IN_OP]
                feet_conf = joints_conf[vi, :, FEET_IN_OP].reshape(-1, 2, 3)
                # TODO: with using POSA
                feet_conf = feet_conf.mean(-1)#[0] before 0613 # B, 2
                reliable_feet = feet_conf >= 0.4
                logger.info(f'reliable_feet length: {reliable_feet.shape}')
                # use scene object contact information
                # contact_flag = scene_model.get_contact_flag(body_vertices_world)

                # import pdb;pdb.set_trace()
                # if the feet far away from the ground plane
                contact_flag = (gt_ground_plane - torch.max(feet_vertices, -1)[0]) > 0.1

                into_gp = (torch.max(feet_vertices, -1)[0] - gt_ground_plane) > 0
                # not reliable feet, max(feet) < gp, if exists contact with other objects, use GM contact loss
                rho_gp_flag = reliable_feet | (~reliable_feet & contact_flag & ~into_gp)
                logger.info(f'rho_gp_flag length: {rho_gp_flag.sum()}')
                if self.rho_gp_flag is None:
                    self.rho_gp_flag = rho_gp_flag

                ############## GM_contact_ground_loss
                # two feets contact gp
                tmp_gp_contact_dist = F.mse_loss(feet_vertices, gt_ground_plane.expand_as(feet_vertices), reduction='none').sqrt()
                tmp_gp_contact_loss = self.gp_contact_robustifier(tmp_gp_contact_dist)
                # logger.info(f'tmp_gp_contact_loss length: {tmp_gp_contact_loss.shape}')
                reliable_gp_contact_loss = tmp_gp_contact_loss.mean(-1) * self.rho_gp_flag.detach().type(torch.float)

                # gp_contact_loss = self.robustifier(smplx_model_vertices[:, ground_contact_vertices_ids, 1] - self.ground_plane)
                # gp_contact_loss = (ground_contact_value * gp_contact_loss.mean(-1)).sum() / ( num_labels + 1e-9)

                ############## abs_loss
                # two feets support by a ground plane
                # not reliable feet, max(feet) > gp
                l1_gp_dist = torch.abs(torch.max(feet_vertices.reshape(batch_size, -1), -1)[0] - gt_ground_plane).unsqueeze(1).expand(-1, 2)
                # not reliable feet, max(feet) < gp, check contact, if not
                l1_gp_flag = (~reliable_feet) & (into_gp | (~contact_flag))
                logger.info(f'l1_gp_flag length: {l1_gp_flag.sum()}')
                if self.l1_gp_flag is None:
                    self.l1_gp_flag = l1_gp_flag
                gp_l1_loss = l1_gp_dist * self.l1_gp_flag.detach().type(torch.float)

                assert (rho_gp_flag | l1_gp_flag).sum() == rho_gp_flag.numel()
                # import pdb;pdb.set_trace()
                # gp_contact_loss = (reliable_gp_contact_loss + gp_l1_loss).mean() # * batch_size

                gp_contact_loss = l1_gp_dist.mean()

                if torch.isnan(gp_contact_loss):
                    import pdb;pdb.set_trace()
        else:
            if self.ground_contact_support and self.gp_contact_loss_weight > 0:
                import pdb;pdb.set_trace()
                gt_ground_plane = scene_model.ground_plane
                feet_vertices = body_vertices_world[:, self.ground_contact_vertices_ids, 1] # B, 2, N
                # reliable feet: op conf > thre
                feet_conf = joints_conf[vi, :, FEET_IN_OP].reshape(-1, 2, 3)
                # TODO: with using POSA
                feet_conf = feet_conf.mean(-1)#[0] before 0613 # B, 2
                reliable_feet = feet_conf >= 0.4
                logger.info(f'reliable_feet length: {reliable_feet.shape}')
                # use scene object contact information
                # contact_flag = scene_model.get_contact_flag(body_vertices_world)

                # if the feet far away from the ground plane
                contact_flag = (gt_ground_plane - torch.max(feet_vertices, -1)[0]) > 0.15

                into_gp = (torch.max(feet_vertices, -1)[0] - gt_ground_plane) > 0
                # not reliable feet, max(feet) < gp, if exists contact with other objects, use GM contact loss
                rho_gp_flag = reliable_feet | (~reliable_feet & contact_flag & ~into_gp)
                logger.info(f'rho_gp_flag length: {rho_gp_flag.sum()}')
                if self.rho_gp_flag is None:
                    self.rho_gp_flag = rho_gp_flag

                ############## GM_contact_ground_loss
                # two feets contact gp
                tmp_gp_contact_dist = F.mse_loss(feet_vertices, gt_ground_plane.expand_as(feet_vertices), reduction='none').sqrt()
                tmp_gp_contact_loss = self.gp_contact_robustifier(tmp_gp_contact_dist)
                # logger.info(f'tmp_gp_contact_loss length: {tmp_gp_contact_loss.shape}')
                reliable_gp_contact_loss = tmp_gp_contact_loss.mean(-1) * self.rho_gp_flag.detach().type(torch.float)

                # gp_contact_loss = self.robustifier(smplx_model_vertices[:, ground_contact_vertices_ids, 1] - self.ground_plane)
                # gp_contact_loss = (ground_contact_value * gp_contact_loss.mean(-1)).sum() / ( num_labels + 1e-9)

                ############## abs_loss
                # two feets support by a ground plane
                # not reliable feet, max(feet) > gp
                l1_gp_dist = torch.abs(torch.max(feet_vertices.reshape(batch_size, -1), -1)[0] - gt_ground_plane).unsqueeze(1).expand(-1, 2)
                # not reliable feet and not far away from gp | max(feet) < gp
                l1_gp_flag = ((~reliable_feet) & (~contact_flag)) | into_gp
                logger.info(f'l1_gp_flag length: {l1_gp_flag.sum()} / {l1_gp_flag.numel()}')
                if self.l1_gp_flag is None:
                    self.l1_gp_flag = l1_gp_flag
                gp_l1_loss = l1_gp_dist * self.l1_gp_flag.detach().type(torch.float)

                # assert (rho_gp_flag | l1_gp_flag).sum() == rho_gp_flag.numel()
                # import pdb;pdb.set_trace()
                # gp_contact_loss = (reliable_gp_contact_loss + gp_l1_loss).mean() # * batch_size
                gp_contact_loss = (gp_l1_loss).mean()
                if torch.isnan(gp_contact_loss):
                    import pdb;pdb.set_trace()

        # print(kwargs['use_scene_loss'])
        if self.use_scene_loss: # add scene constraints from PROX / POSA
            if False:
                ### scene constraint loss.
                sdf_penetration_loss = 0.0
                # TODO: check whether is batch size average?
                # TODO: add scene sdf calculation in scene_model.
                if self.sdf_penetration and self.sdf_penetration_loss_weight >0:
                    assert scene_model is not None
                    cnt = 0
                    for idx in range(batch_size):
                        # import pdb;pdb.set_trace()
                        sdf_penetration_loss_tmp = scene_model(body_vertices_world[idx:idx+1], stage=1)
                        if sdf_penetration_loss_tmp > 0:
                            cnt = cnt + 1
                        sdf_penetration_loss = sdf_penetration_loss + sdf_penetration_loss_tmp
                        #stage=11) not reasonable
                    # sdf_penetration_loss = self.sdf_penetration_loss_weight * sdf_penetration_loss / batch_size
                    sdf_penetration_loss = self.sdf_penetration_loss_weight * sdf_penetration_loss / (cnt+1e-9)

                contact_loss = 0.0
                # TODO: add scene v, vn calculation in scene model.
                if self.contact and self.contact_loss_weight >0:
                    # import pdb;pdb.set_trace()
                    assert scene_model is not None
                    cnt = 0
                    for idx in range(batch_size):
                        # if idx == 7:
                        #     import pdb;pdb.set_trace()
                        tmp_contact = scene_model(body_vertices_world[idx:idx+1], stage=2, contact_verts_ids=self.contact_verts_ids,
                                                contact_angle=self.contact_angle, contact_robustifier=self.contact_robustifier, ftov=self.ftov)
                        # logger.info(f'{idx}: {tmp_contact.item()}')

                        if tmp_contact > 0:
                            cnt = cnt + 1
                        contact_loss = contact_loss + tmp_contact
                    # contact_loss = self.contact_loss_weight * contact_loss / batch_size
                    contact_loss = self.contact_loss_weight * contact_loss / (cnt+1e-9)

                # TODO: calculate depth map for scene model, and use it as constraint.
                depth_loss = 0.0
                if self.ordinal_depth and self.ordinal_depth_loss_weight > 0:
                    cnt = 0
                    assert scene_model is not None
                    for idx in range(batch_size):
                    # input current new per-frame mask.
                        # import pdb;pdb.set_trace()
                        # this
                        tmp_depth_loss = scene_model.compute_ordinal_depth_loss_perframe(body_vertices_world[idx:idx+1], idx, use_for_human=self.use_human_depth)['loss_depth']
                        # import pdb;pdb.set_trace()
                        if torch.isnan(tmp_depth_loss): # the depth loss is too big, generates large gradient, leads to error !!!
                            import pdb;pdb.set_trace()
                            # tmp_depth_loss = scene_model.compute_ordinal_depth_loss_perframe(body_vertices_world[idx:idx+1], idx, use_for_human=self.use_human_depth)['loss_depth']

                        if tmp_depth_loss > 0:
                            cnt += 1

                        depth_loss = depth_loss + self.ordinal_depth_loss_weight * tmp_depth_loss
                    # depth_loss = depth_loss / batch_size
                    depth_loss = depth_loss / (cnt+1e-9)
                    # ! Warning: add robustifier: for Nan problem, the gradients are so big! on 03.08
                    # import pdb;pdb.set_trace()
                    # if self.depth_robustifier is not None:
                    #     squared_res = depth_loss ** 2
                    #     depth_loss = torch.div(squared_res, squared_res + self.depth_robustifier ** 2)
            else:
                # loss weight > 0: leads to results;
                scene_constraint_loss_weight = {}
                sdf_penetration_loss = torch.tensor(0.0).cuda()
                contact_loss = torch.tensor(0.0).cuda()
                depth_loss = torch.tensor(0.0).cuda()
                if self.sdf_penetration and self.sdf_penetration_loss_weight >0:
                    scene_constraint_loss_weight['lw_sdf'] = self.sdf_penetration_loss_weight
                if self.contact and self.contact_loss_weight >0:
                    scene_constraint_loss_weight['lw_contact'] = self.contact_loss_weight
                # if self.posa and self.posa_loss_weight >0:
                #     scene_constraint_loss_weight['lw_posa'] = self.posa_loss_weight
                #     posa
                if self.ordinal_depth and self.ordinal_depth_loss_weight > 0:
                    scene_constraint_loss_weight['lw_depth'] = self.ordinal_depth_loss_weight
                scene_loss_dict = scene_model(body_vertices_world,
                                                save_dir=self.save_dir,
                                                loss_weights=scene_constraint_loss_weight,
                                                stage=3, contact_verts_ids=self.contact_verts_ids,
                                                contact_angle=self.contact_angle, \
                                                contact_robustifier=self.gp_contact_robustifier, ftov=self.ftov) #self.contact_robustifier
                if 'loss_sdf' in scene_loss_dict:
                    sdf_penetration_loss = scene_loss_dict['loss_sdf'] * self.sdf_penetration_loss_weight
                if 'loss_contact' in scene_loss_dict:
                    contact_loss = scene_loss_dict['loss_contact'] * self.contact_loss_weight
                if 'loss_depth' in scene_loss_dict:
                    depth_loss = scene_loss_dict['loss_depth'] * self.ordinal_depth_loss_weight
                if 'loss_posa' in scene_loss_dict:
                    contact_loss = scene_loss_dict['loss_contact'] * self.posa_loss_weight

                # best version for only contact & sdf loss: 1, 1
        else:
            sdf_penetration_loss = 0.0
            contact_loss = 0.0
            depth_loss = 0.0
        scene_loss = 0.0
        message_scene = ''
        # if self.scene and self.scene_loss_weight > 0:
        #     assert scene_model  is not None
        #     init_pid = 0
        #     scene_local_loss_weights = DEFAULT_LOSS_WEIGHTS[f'stage2_init{init_pid}']['loss_weight']
        #     scene_loss_tmp = scene_model(stage=0, loss_weights=scene_local_loss_weights)
        #     scene_loss_dict_weighted = {
        #             k: scene_loss_tmp[k] * scene_local_loss_weights[k.replace("loss", "lw")] for k in scene_loss_tmp
        #         }

        #     for key, val in scene_loss_tmp.items():
        #             message_scene += f' {key}: {val.item()}_{scene_local_loss_weights[key.replace("loss", "lw")]}_{scene_loss_dict_weighted[key].item()}'
        #     scene_loss = sum(scene_loss_dict_weighted.values()).sum()
        #     scene_loss = self.scene_loss_weight * scene_loss

        # TODO: experiments to analysis;
        message = f'ST2 joint: {joint_loss}, \
            angle_prior: {angle_prior_loss}, pen:{pen_loss}, pare_prior:{pare_pose_prior_loss}, \
            gp: {gp_support_loss}, \
            gp_contact: {gp_contact_loss}, sdf: {sdf_penetration_loss}, contact: {contact_loss}, \
            scene: {scene_loss}, depth: {depth_loss}, smooth_j2d_loss:{smooth_j2d_loss}, smooth_j3d_loss: {smooth_j3d_loss}, motion_smooth_prior: {motion_prior_loss}'
        message += ' || scene details: ' + message_scene
        logger.info(message)

        if self.loss_use_sum:
            total_loss = (joint_loss + pprior_loss + shape_loss +
                        angle_prior_loss + pen_loss + pare_pose_prior_loss +
                        jaw_prior_loss + expression_loss +
                        left_hand_prior_loss + right_hand_prior_loss
                        + gp_support_loss + gp_contact_loss +
                        sdf_penetration_loss + contact_loss + depth_loss + smooth_j2d_loss + smooth_j3d_loss + motion_prior_loss) * batch_size + scene_loss
        else:
            total_loss = (joint_loss + pprior_loss + shape_loss +
                        angle_prior_loss + pen_loss + pare_pose_prior_loss +
                        jaw_prior_loss + expression_loss +
                        left_hand_prior_loss + right_hand_prior_loss
                        + gp_support_loss + gp_contact_loss +
                        sdf_penetration_loss + contact_loss + depth_loss + smooth_j2d_loss + smooth_j3d_loss + motion_prior_loss) + scene_loss

        if 'tb_debug' in kwargs and kwargs['tb_debug']:
            debug_loss_dict = {
                'loss_joint': joint_loss,
                'loss_pprior': pprior_loss,
                'loss_shape': shape_loss,
                'loss_angle_prior': angle_prior_loss,
                'loss_pare_pose_prior': pare_pose_prior_loss,
                'loss_expression': expression_loss,
                'loss_left_hand_prior': left_hand_prior_loss,
                'loss_right_hand_prior': right_hand_prior_loss,
                'loss_gp_support_body': gp_support_loss,
                'loss_gp_contact':  gp_contact_loss,
                'loss_sdf_penetration': sdf_penetration_loss,
                'loss_contact': contact_loss,
                'loss_scene': scene_loss,
                'loss_depth': depth_loss,
                'smooth_j3d_loss': smooth_j3d_loss,
                'smooth_j2d_loss': smooth_j2d_loss,
                'motion_prior_loss': motion_prior_loss,
                'loss_total_MultiViewSMPLify': total_loss,
            }

            if self.ground_plane_support:
                debug_loss_dict['SMPLX_weight_gp_support'] = self.gp_support_loss_weight
            if self.ground_contact_support:
                debug_loss_dict['SMPLX_weight_gp_contact'] = self.gp_contact_loss_weight
            if self.sdf_penetration:
                debug_loss_dict['SMPLX_weight_sdf_penetration'] = self.sdf_penetration_loss_weight
            if self.contact:
                debug_loss_dict['SMPLX_weight_contact'] = self.contact_loss_weight
            if self.scene:
                debug_loss_dict['SMPLX_weight_scene'] = self.scene_loss_weight
            if self.ordinal_depth:
                debug_loss_dict['SMPLX_weight_ordinal_depth'] = self.ordinal_depth_loss_weight

            if self.scene and self.scene_loss_weight > 0:
                for key, value in scene_local_loss_weights.items():
                    debug_loss_dict[f'scene_weight_{key}'] = value
                for key in scene_loss_tmp.keys():
                    debug_loss_dict[f'loss_scene_{key}_wt'] = scene_loss_dict_weighted[key]
                    debug_loss_dict[f'loss_scene_{key}'] = scene_loss_tmp[key]

            if self.video_smooth:
                debug_loss_dict['smooth_2d_weight'] = self.smooth_2d_weight
                debug_loss_dict['smooth_3d_weight'] = self.smooth_3d_weight
                if self.motion_smooth_prior_flag:
                    debug_loss_dict['motion_prior_weight'] = self.motion_prior_weight

            return total_loss, debug_loss_dict
        else:
            return total_loss


import pygame
import random

pygame.init()
window = pygame.display.set_mode((1920, 1080))
pygame.display.set_caption("ULTIMATE GHOST FIGHT 9000")

class player:
    def __init__(
        self,
        x,
        y,
        width,
        height,
        ghostLeft,
        ghostRight,
        ghostCrouchLeft,
        ghostCrouchRight,
        ball,
        left,
        right,
        balls,
        ghostShootLeft,
        ghostShootRight,
    ):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.velocity = 6
        self.jumps = False
        self.jumpPhase = -25
        self.crouch = False
        self.right = right
        self.left = left
        self.ghostLeft = ghostLeft
        self.ghostRight = ghostRight
        self.ghostCrouchLeft = ghostCrouchLeft
        self.ghostCrouchRight = ghostCrouchRight
        self.ball = ball
        self.balls = balls
        self.health = 10
        self.shootPhase = -1
        self.lastShot = 0
        self.ballLoad = 5
        self.ballReload = 0
        self.ghostShootLeft = ghostShootLeft
        self.ghostShootRight = ghostShootRight
        self.gravity = 5
        self.floatPhase = 0
        self.float = 0
        self.score = 0
        self.crouchTime = 3

def objectRender(objects):
    for object in objects:
        window.blit(object.image, (object.x, object.y))
    window.blit(
        pygame.font.SysFont("comicsansms", 100).render(
            str(player1.score), 1, (255, 180, 22)
        ),
        (30, 0),
    )
    window.blit(
        pygame.font.SysFont("comicsansms", 100).render(
            str(player2.score), 1, (255, 180, 22)
        ),
        (
            1890
            - (
                pygame.font.SysFont("comicsansms", 100)
                .render(str(player2.score), 1, (255, 180, 22))
                .get_width()
            ),
            0,
        ),
    )

def movementDrawShoot(player, left, right, jump, crouch, shoot, objects):
    if player.ballLoad < 5:
        player.ballReload += 1
        if player.ballReload % 75 == 0:
            player.ballLoad += 1
    else:
        player.ballReload = 0
    global hearts
    for heart in hearts:
        if (
            heart.x <= player.x + player.width
            and heart.x + heart.width >= player.x
            and player.health < 10
            and heart.y >= player.y
            and heart.y <= player.y + player.height
        ):
            hearts.pop(hearts.index(heart))
            if player.health > 5:
                player.health += 10 - player.health
            else:
                player.health += 5
    if player.crouchTime > 3 and player.crouch is False:
        player.crouchTime -= 1
    player.floatPhase += 1
    if player.floatPhase % 30 == 0 and player.float == 0 and player.jumps is False:
        player.y += 8
        player.float = 1
    elif player.floatPhase % 30 == 0 and player.float == 1 and player.jumps is False:
        player.y -= 8
        player.float = 0
    player.y += player.gravity
    player.gravity = 5
    for object in objects:
        if (
            player.y + player.height >= object.y
            and player.y + player.height <= object.y + object.height / 5
            and player.x > object.x - player.width / 2
            and player.x < object.x + object.width - player.width / 2
        ):
            player.gravity = 0
            if player.jumps and player.y + player.height > object.y:
                player.y = object.y - player.height
                player.jumps = False
                player.jumpPhase = -25
    for ball in player.balls:
        if ball.x > 0 and ball.x < 1920:
            ball.x += ball.velocity
        else:
            player.balls.pop(player.balls.index(ball))
    if player.y > 1080:
        player.health = 0
    if player.health > 0:
        if shoot and player.crouch is False and player.ballLoad > 0:
            player.shootPhase += 1
            if player.left is True and player.shootPhase % 15 == 0:
                player.ballLoad -= 1
                facing = -1
                player.balls.append(
                    projectile(player.x + 56, player.y + 28, 12, player.ball, facing)
                )
            elif player.right is True and player.shootPhase % 15 == 0:
                player.ballLoad -= 1
                facing = 1
                player.balls.append(
                    projectile(player.x + 56, player.y + 28, 12, player.ball, facing)
                )
        elif player.shootPhase != -1:
            player.lastShot += 1
            if player.lastShot == 15:
                player.shootPhase == -1
        if left and player.x > 0:
            player.x -= player.velocity
            player.left = True
            player.right = False
        if right and player.x + player.width < 1920:
            player.x += player.velocity
            player.right = True
            player.left = False
        if not (player.jumps):
            if crouch and player.crouch is False and player.crouchTime == 3:
                player.y = player.y + 50
                player.height = 70
                player.crouch = True
            if not crouch and player.crouch is True or player.crouchTime >= 296:
                player.y = player.y - 50
                player.height = 120
                player.crouch = False
            for object in objects:
                if (
                    jump
                    and player.crouch is False
                    and player.y + player.height >= object.y
                    and player.y + player.height <= object.y + object.height / 5
                    and player.x > object.x - player.width / 2
                    and player.x < object.x + object.width - player.width / 2
                ):
                    player.jumps = True
        else:
            if player.jumpPhase <= 25:
                player.y += player.jumpPhase
                player.jumpPhase += 1
            else:
                player.jumps = False
                player.jumpPhase = -25
        if shoot and not crouch and player.ballLoad > 0:
            if player.left is True and left and not right:
                window.blit(player.ghostShootLeft, (player.x, player.y))
            elif player.right is True and right and not left:
                window.blit(player.ghostShootRight, (player.x, player.y))
            elif player.left is True:
                window.blit(player.ghostShootLeft, (player.x, player.y))
            elif player.right is True:
                window.blit(player.ghostShootRight, (player.x, player.y))
        elif player.crouch is True:
            if player.left is True and left and not right:
                player.crouchTime += 1
                window.blit(player.ghostCrouchLeft, (player.x, player.y))
            elif player.right is True and right and not left:
                player.crouchTime += 1
                window.blit(player.ghostCrouchRight, (player.x, player.y))
            elif player.left is True:
                player.crouchTime += 1
                window.blit(player.ghostCrouchLeft, (player.x, player.y))
            elif player.right is True:
                player.crouchTime += 1
                window.blit(player.ghostCrouchRight, (player.x, player.y))
        else:
            if player.left is True and left and not right:
                window.blit(player.ghostLeft, (player.x, player.y))
            elif player.right is True and right and not left:
                window.blit(player.ghostRight, (player.x, player.y))
            elif player.left is True:
                window.blit(player.ghostLeft, (player.x, player.y))
            elif player.right is True:
                window.blit(player.ghostRight, (player.x, player.y))
        for ball in player.balls:
            window.blit(ball.image, (ball.x, ball.y))
        pygame.draw.rect(window, (255, 255, 255), (player.x, player.y - 28, 100, 19))
        pygame.draw.rect(
            window,
            (255, 0, 0),
            (player.x + 1, player.y - 18, 10 * player.health - 2, 8),
        )
        pygame.draw.rect(
            window,
            (0, 0, 255),
            (player.x + 1, player.y - 27, (299 - player.crouchTime) // 3, 8),
        )
        for i in range(0, player.ballLoad):
            pygame.draw.rect(
                window, (255, 255, 255), (player.x + i * 10, player.y - 40, 8, 8)
            )
    else:
        if player1.health == 0:
            player2.score += 1
            window.blit(
                pygame.font.SysFont("comicsansms", 100).render(
                    "black ghost wins", 1, (255, 180, 22)
                ),
                (
                    960
                    - pygame.font.SysFont("comicsansms", 100)
                    .render("black ghost wins", 1, (255, 180, 22))
                    .get_width()
                    / 2,
                    540
                    - pygame.font.SysFont("comicsansms", 100)
                    .render("black ghost wins", 1, (255, 180, 22))
                    .get_height()
                    / 2,
                ),
            )
        elif player2.health == 0:
            player1.score += 1
            window.blit(
                pygame.font.SysFont("comicsansms", 100).render(
                    "white ghost wins", 1, (255, 180, 22)
                ),
                (
                    960
                    - pygame.font.SysFont("comicsansms", 100)
                    .render("white ghost wins", 1, (255, 180, 22))
                    .get_width()
                    / 2,
                    540
                    - pygame.font.SysFont("comicsansms", 100)
                    .render("white ghost wins", 1, (255, 180, 22))
                    .get_height()
                    / 2,
                ),
            )
        pygame.mixer.music.stop()
        music = pygame.mixer.music.load("media/music.mp3")
        pygame.mixer.music.play()
        pygame.display.update()
        player1.x = 150
        player1.y = 400
        player2.x = 1670
        player2.y = 400
        player1.left = False
        player1.right = True
        player2.left = True
        player2.right = False
        player1.health = 10
        player2.health = 10
        hearts = []
        player1.balls = []
        player2.balls = []
        player1.ballLoad = 5
        player1.ballReload = 0
        player1.lastShot = 0
        player2.ballLoad = 5
        player2.ballReload = 0
        player2.lastShot = 0
        lifeCount = 1
        pygame.time.delay(2000)

def hit(player, enemy, objects):
    for ball in player.balls:
        for object in objects:
            if (
                ball.y - ball.radius < object.y + object.height
                and ball.y + ball.radius > object.y
                and ball.x + ball.radius > object.x
                and ball.x - ball.radius < object.x + object.width
            ):
                player.balls.pop(player.balls.index(ball))
        if (
            ball.y - ball.radius < enemy.y + enemy.height
            and ball.y + ball.radius > enemy.y
        ):
            if (
                ball.x + ball.radius > enemy.x
                and ball.x - ball.radius < enemy.x + enemy.width
            ):
                player.balls.pop(player.balls.index(ball))
                enemy.health -= 1

class projectile:
    def __init__(self, x, y, radius, image, facing):
        self.x = x
        self.y = y
        self.radius = radius
        self.image = image
        self.facing = facing
        self.velocity = 8 * facing

class object:
    def __init__(self, x, y, width, height, image):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.image = image

class life:
    def __init__(self, x):
        self.x = x
        self.y = 494
        self.width = 36
        self.height = 36
        self.image = pygame.image.load("media/life.png").convert_alpha()
        self.floatPhase = 0
        self.float = 0

background = pygame.image.load("media/background.jpg").convert_alpha()

class title:
    def __init__(self, image):
        self.image = image
        self.x = background.get_width() / 2 - self.image.get_width() / 2
        self.y = 340
        self.float = 0
        self.floatPhase = 0

def lifeGeneration():
    global lifeCount
    lifeCount += 1
    if lifeCount % 700 == 0:
        heart = life(random.choice([182, 1702]))
        hearts.append(heart)
        if len(hearts) > 1:
            hearts.pop(hearts.index(hearts[0]))
    for heart in hearts:
        heart.floatPhase += 1
        if heart.floatPhase % 20 == 0 and heart.float == 0:
            heart.y -= 8
            heart.float = 1
        elif heart.floatPhase % 20 == 0 and heart.float == 1:
            heart.y += 8
            heart.float = 0
        window.blit(heart.image, (heart.x, heart.y))

def titleRender(title):
    title.floatPhase += 1
    if title.floatPhase % 80 == 0 and title.float == 0:
        title.y -= 8
        title.float = 1
    elif title.floatPhase % 80 == 0 and title.float == 1:
        title.y += 8
        title.float = 0
    window.blit(title.image, (title.x, title.y))

title = title(pygame.image.load("media/title.png").convert_alpha())
music = pygame.mixer.music.load("media/music.mp3")
player1 = player(
    150,
    400,
    108,
    120,
    pygame.image.load("media/whiteghostleft.png").convert_alpha(),
    pygame.image.load("media/whiteghostright.png").convert_alpha(),
    pygame.image.load("media/whiteghostcrouchleft.png").convert_alpha(),
    pygame.image.load("media/whiteghostcrouchright.png").convert_alpha(),
    pygame.image.load("media/whiteball.png").convert_alpha(),
    False,
    True,
    [],
    pygame.image.load("media/whiteghostshootleft.png").convert_alpha(),
    pygame.image.load("media/whiteghostshootright.png").convert_alpha(),
)
player2 = player(
    1670,
    400,
    108,
    120,
    pygame.image.load("media/blackghostleft.png").convert_alpha(),
    pygame.image.load("media/blackghostright.png").convert_alpha(),
    pygame.image.load("media/blackghostcrouchleft.png").convert_alpha(),
    pygame.image.load("media/blackghostcrouchright.png").convert_alpha(),
    pygame.image.load("media/blackball.png").convert_alpha(),
    True,
    False,
    [],
    pygame.image.load("media/blackghostshootleft.png").convert_alpha(),
    pygame.image.load("media/blackghostshootright.png").convert_alpha(),
)
objects = []
platform1 = object(
    100, 540, 200, 100, pygame.image.load("media/cloud.png").convert_alpha()
)
objects.append(platform1)
platform2 = object(
    1620, 540, 200, 100, pygame.image.load("media/cloud.png").convert_alpha()
)
objects.append(platform2)
platform3 = object(
    450, 540, 1020, 100, pygame.image.load("media/longcloud.png").convert_alpha()
)
objects.append(platform3)
platform4 = object(
    660, 340, 600, 90, pygame.image.load("media/shortcloud.png").convert_alpha()
)
objects.append(platform4)
lifeCount = 1
hearts = []
menu = True
while menu:
    pygame.time.delay(5)
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            menu = False
            play = False
    window.fill((0, 0, 0))
    titleRender(title)
    keys = pygame.key.get_pressed()
    if keys[pygame.K_RETURN]:
        menu = False
        play = True
    pygame.display.update()
pygame.mixer.music.play()
while play:
    pygame.time.delay(5)
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            play = False
    window.blit(background, (0, 0))
    keys = pygame.key.get_pressed()
    objectRender(objects)
    lifeGeneration()
    movementDrawShoot(
        player1,
        keys[pygame.K_a],
        keys[pygame.K_d],
        keys[pygame.K_w],
        keys[pygame.K_s],
        keys[pygame.K_SPACE],
        objects,
    )
    movementDrawShoot(
        player2,
        keys[pygame.K_LEFT],
        keys[pygame.K_RIGHT],
        keys[pygame.K_UP],
        keys[pygame.K_DOWN],
        keys[pygame.K_KP0],
        objects,
    )
    hit(player1, player2, objects)
    hit(player2, player1, objects)
    pygame.display.update()
pygame.quit()


import json
import re
from typing import Generator

import milagro_bls_binding as bls
import pytest
from aioresponses import CallbackResult, aioresponses
from eth_typing import HexStr
from staking_deposit.key_handling.keystore import Keystore
from web3 import Web3

from src.validators.typings import BLSPrivkey

@pytest.fixture
def remote_signer_url() -> str:
    return 'http://web3signer:9000'

@pytest.fixture
def mocked_remote_signer(
    remote_signer_url: str,
) -> Generator:
    _remote_signer_pubkey_privkey_mapping: dict[HexStr, BLSPrivkey] = {}

    def _mocked_keymanager_import_endpoint(url, **kwargs) -> CallbackResult:
        data = kwargs['json']
        keystores = [Keystore.from_json(json.loads(keystore)) for keystore in data['keystores']]
        passwords = data['passwords']

        for keystore, password in zip(keystores, passwords):
            priv_key = BLSPrivkey(keystore.decrypt(password))
            pubkey = Web3.to_hex(bls.SkToPk(priv_key))
            _remote_signer_pubkey_privkey_mapping[pubkey] = priv_key

        return CallbackResult(
            status=200, payload={'data': [{'status': 'imported'} for _ in keystores]}
        )

    def _mocked_keymanager_delete_endpoint(url, **kwargs) -> CallbackResult:
        data = kwargs['json']
        pubkeys = data['pubkeys']

        for pubkey in pubkeys:
            _remote_signer_pubkey_privkey_mapping.pop(pubkey)

        return CallbackResult(
            status=200, payload={'data': [{'status': 'deleted'} for _ in pubkeys]}
        )

    def _mocked_keymanager_list_endpoint(url, **kwargs) -> CallbackResult:
        return CallbackResult(
            status=200,
            payload={
                'data': [
                    {'validating_pubkey': pubkey}
                    for pubkey in _remote_signer_pubkey_privkey_mapping.keys()
                ]
            },
        )

    def _mocked_sign_endpoint(url, **kwargs) -> CallbackResult:
        public_key_to_sign_for = url.path.split('/')[-1]

        try:
            corresponding_private_key = _remote_signer_pubkey_privkey_mapping[
                public_key_to_sign_for
            ]
        except KeyError:
            return CallbackResult(status=404, body='Not Found')

        signature = bls.Sign(
            corresponding_private_key, Web3.to_bytes(hexstr=kwargs['json']['signing_root'])
        )

        return CallbackResult(payload={'signature': f'0x{signature.hex()}'})

    with aioresponses() as m:
        # Mocked keymanager list keys endpoint
        m.get(
            f'{remote_signer_url}/eth/v1/keystores',
            callback=_mocked_keymanager_list_endpoint,
            repeat=True,
        )

        # Mocked keymanager import keystores endpoint
        m.post(
            f'{remote_signer_url}/eth/v1/keystores',
            callback=_mocked_keymanager_import_endpoint,
            repeat=True,
        )

        # Mocked keymanager delete keys endpoint
        m.delete(
            f'{remote_signer_url}/eth/v1/keystores',
            callback=_mocked_keymanager_delete_endpoint,
            repeat=True,
        )

        # Mocked signing endpoint
        m.post(
            re.compile(f'^{remote_signer_url}/api/v1/eth2/sign/\\w{{98}}$'),
            callback=_mocked_sign_endpoint,
            repeat=True,
        )
        yield


from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.firefox.options import Options

options = Options()
options.add_argument('--headless')

firefox = webdriver.Firefox(options=options)

URL = 'https://diolinux.com.br/'

def scrape(url):
    firefox.get(url)
    posts = []

    for post in firefox.find_elements(By.CLASS_NAME, 'inhype-grid-post'):
        new_item = dict()

        new_item['title'] = post.find_element(
            By.CLASS_NAME, 'post-title',
        ).find_element(By.TAG_NAME, 'a').get_attribute('innerHTML')

        new_item['link'] = post.find_element(
            By.CLASS_NAME, 'post-title',
        ).find_element(By.TAG_NAME, 'a').get_attribute('href')

        posts.append(new_item)
    return posts

posts = scrape('https://diolinux.com.br/')

for post in posts:
    print(post)


import os
import time

import psutil

import Natsunagi.modules.no_sql.users_db as users_db
from Natsunagi import StartTime
from Natsunagi.modules.helper_funcs import formatter

# Stats Module

async def bot_sys_stats():
    bot_uptime = int(time.time() - StartTime)
    cpu = psutil.cpu_percent()
    mem = psutil.virtual_memory().percent
    disk = psutil.disk_usage("/").percent
    process = psutil.Process(os.getpid())
    stats = f"""
arya@FurryChemistry:~$ Natsunagi:
------------------
Natsunagi Nagisa Uptime: {formatter.get_readable_time((bot_uptime))}
Bot Capasity: {round(process.memory_info()[0] / 1024 ** 2)} MB
CPU Usage: {cpu}%
RAM Usage: {mem}%
Disk Usage: {disk}%
Users: {users_db.num_users()} users.
Groups: {users_db.num_chats()} groups.
"""

    return stats


import pygame
import random
from os import path

img_dir = path.join(path.dirname(__file__), 'img3')
snd_dir = path.join(path.dirname(__file__), 'snd3')

pygame.display.set_caption('Ближе-дальше')

FPS = 60
WIDTH = 800
HEIGHT = 600

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (125, 125, 125)
LIGHT_BLUE = (64, 128, 255)
GREEN = (0, 200, 64)
YELLOW = (225, 225, 0)
PINK = (230, 50, 230)

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()

corsack_image = pygame.image.load(path.join(img_dir, "корсак3.png")).convert()
corsack = corsack_image.get_rect()
corsack.x = random.randrange(50, 600)
corsack.y = random.randrange(375, 400)
corsack_image.set_colorkey(WHITE)

lion_image = pygame.image.load(path.join(img_dir, "лев4.png")).convert()
lion = lion_image.get_rect()
lion.x = random.randrange(50, 600)
lion.y = random.randrange(450, 500)
lion_image.set_colorkey(WHITE)

nadpis = pygame.image.load(path.join(img_dir, "Кто дальше2.png")).convert()
nadpis.set_colorkey(BLACK)
nadpis_Rect = nadpis.get_rect()
nadpis_Rect.topleft = (280, 30)

music_image = pygame.image.load(path.join(img_dir, "музыка2.png"))
music_button = music_image.get_rect()
music_button.topleft = (0, 550)
music_image.set_colorkey(BLACK)

background = pygame.image.load(path.join(img_dir, "пустыня2.jpg")).convert()
background_rect = background.get_rect()

pygame.mixer.music.load(path.join(snd_dir, 'Desert.mp3'))
pygame.mixer.music.set_volume(0.4)

incorrect_sound = pygame.mixer.Sound(path.join(snd_dir, 'звук неправильного ответа.mp3'))
incorrect_sound.set_volume(0.35)

kto_dalshe_sound = pygame.mixer.Sound(path.join(snd_dir,'Кто дальше.mp3'))
kto_dalshe_sound.play()

question_image = pygame.image.load(path.join(img_dir,'вопрос2.png')).convert()
question_button = question_image.get_rect()
question_button.topleft = (750,550)
question_image.set_colorkey(BLACK)

pygame.mixer.music.play()

pygame.mixer.music.play(loops=-1)

running = True
flPause = False
while running:
    clock.tick(FPS)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        if event.type == pygame.MOUSEBUTTONDOWN:
            pos = pygame.mouse.get_pos()
            if corsack.collidepoint(pos) and corsack.y < lion.y:
                print("Получилось")
                pygame.mixer.music.stop()
                kto_dalshe_sound.stop()
                sc = pygame.display.set_mode((500, 300))
                clock = pygame.time.Clock()
                start_time = pygame.time.get_ticks()
                f = pygame.font.Font(None, 100)

                molodec_image = pygame.image.load(path.join(img_dir, "Смайлик2.png"))
                molodec = molodec_image.get_rect()
                molodec.topleft = (150, 50)
                molodec_sound = pygame.mixer.Sound(path.join(snd_dir, 'Молодец.mp3'))
                molodec_sound.play()

                sc.blit(molodec_image, molodec)
                r = True
                while r == True:
                    sc.fill((255, 255, 255))
                    for i in pygame.event.get():
                        if i.type == pygame.QUIT:
                            exit()

                    clock.tick(FPS)

                    if start_time and pygame.time.get_ticks() - start_time > 2000:
                        r = False
                    sc.blit(molodec_image, molodec)
                    pygame.display.update()
                running = False
                import Diplom.Main_Menu
            elif lion.collidepoint(pos) and lion.y < corsack.y:
                print("Получилось")
                pygame.mixer.music.stop()
                kto_dalshe_sound.stop()
                sc = pygame.display.set_mode((500, 300))
                clock = pygame.time.Clock()
                start_time = pygame.time.get_ticks()
                f = pygame.font.Font(None, 100)

                molodec_image = pygame.image.load(path.join(img_dir, "Смайлик2.png"))
                molodec = molodec_image.get_rect()
                molodec.topleft = (150, 50)
                molodec_sound = pygame.mixer.Sound(path.join(snd_dir, 'Молодец.mp3'))
                molodec_sound.play()

                sc.blit(molodec_image, molodec)
                r = True
                while r == True:
                    sc.fill((255, 255, 255))
                    for i in pygame.event.get():
                        if i.type == pygame.QUIT:
                            exit()

                    clock.tick(FPS)

                    if start_time and pygame.time.get_ticks() - start_time > 2000:
                        r = False
                    sc.blit(molodec_image, molodec)
                    pygame.display.update()
                running = False
                import Diplom.Main_Menu
            else:
                incorrect_sound.play()
            if music_button.collidepoint(pos):
                incorrect_sound.stop()
                flPause = not flPause
                if flPause:
                    pygame.mixer.music.pause()
                else:
                    pygame.mixer.music.unpause()

            if question_button.collidepoint(pos):
                kto_dalshe_sound.stop()
                incorrect_sound.stop()
                kto_dalshe_sound.play()

        if lion.right > corsack.left and \
                lion.left < corsack.right and \
                lion.bottom > corsack.top and \
                lion.top < corsack.bottom:
            lion.x = random.randrange(50, 600)
            lion.y = random.randrange(375, 400)
            corsack.x = random.randrange(50, 600)
            corsack.y = random.randrange(450, 500)

    screen.blit(background, background_rect)

    screen.blit(corsack_image, corsack)
    screen.blit(lion_image, lion)
    screen.blit(nadpis, nadpis_Rect)
    screen.blit(music_image, music_button)
    screen.blit(question_image, question_button)

    pygame.display.flip()

pygame.quit()


#!/usr/bin/env python3
import pandas as pd
from collections import namedtuple, Counter
import os
from invoke import run
import os.path
import tempfile
from io import StringIO
from bs4 import BeautifulSoup

class AssemblySide(object):
    Top = 0
    Bottom = 1
    Both = 2

AssemblyPosition = namedtuple("AssemblyPosition", ["x", "y", "rot", "side"])

class MPNMissingError(object):
    def __init__(self, refdes):
        self.refdes = refdes

    def __str__(self):
        return f"MPN missing for {self.refdes}"

class PositionMissingError(object):
    def __init__(self, refdes):
        self.refdes = refdes

    def __str__(self):
        return f"Position missing for {self.refdes}"

class Component(object):
    def __init__(self, refdes, value, mpn, footprint, populate=True, position: AssemblyPosition = None, extra_properties={}):
        self.refdes = refdes
        self.value = value
        self.mpn = mpn
        self.footprint = footprint
        self.populate = populate
        self.position = position
        self.extra_properties = extra_properties

    @property
    def errors(self):
        if self.mpn is None:
            self.errors.append(MPNMissingError(self.refdes))
        if self.position is None:
            self.errors.append(PositionMissingError(self.refdes))

    @property
    def mpn_or_value(self):
        if self.mpn is None or self.mpn in ["", "NA", "-", "~"]:
            return self.value
        return self.mpn

    @property
    def mpn_or_value_plus_footprint(self):
        """
        Returns the MPN if it exists, otherwise returns the value and footprint concatenated with a slash.
        """
        if self.mpn is None or self.mpn in ["", "NA", "-", "~"]:
            return f"{self.value} / {self.footprint}"
        return self.mpn

    def asdict(self):
        ret = {
            "Ref": self.refdes,
            "value": self.value,
            "MPN": self.mpn,
            "footprint": self.footprint,
            "populate": self.populate,
        }
        if self.position is not None:
            # Merge dict with dict(self.position)
            ret.update({
                "PosX": self.position.PosX,
                "PosY": self.position.PosY,
                "Rot": self.position.Rot,
                "Side": self.position.Side,
            })
        ret.update(self.extra_properties)
        return ret

    def __str__(self):
        return f"Component({self.refdes})"

    def __repr__(self) -> str:
        return self.__str__()

def extract_components_from_bom(bom, pnp_positions, extra_properties=[]):
    components = []
    for comp in bom.components.find_all("comp"):
        # Find whether to populate or not (<property name="dnp"/>)
        populate: bool = comp.find("property", {"name": "dnp"}) == None
        # Extract RefDes e.g. <comp ref="BC1">
        refdes: str = comp["ref"]
        # Extract value
        value = text_or_None(comp.find("value"))
        # Extract MPN
        mpn = text_or_None(comp.find("field", {"name": "MPN"}))
        # Extract additional properties
        component_extra_properties = {
            property: value_or_None(comp.find("property", {"name": property}))
            for property in extra_properties
        }
        # Filter None values from extra properties
        component_extra_properties = {
            property: value
            for property, value in component_extra_properties.items()
            if value is not None
        }
        # Extract footprint e.g. <footprint>KKS-Microcontroller-Board:10x10mm Laser Data Matrix</footprint>
        footprint = comp.find("footprint").text
        footprint_lib, _ , footprint_name = footprint.partition(":")
        # Extract positions from PNP file
        try:
            position = pnp_positions.loc[refdes]
        except KeyError:
            position = None
        # Add to component list
        components.append(
            Component(
                refdes=refdes,
                value=value,
                mpn=mpn,
                footprint=footprint_name,
                populate=populate,
                position=position,
                extra_properties=component_extra_properties
            )
        )
    return components

def read_kicad_pos_file(filename_or_file):
    pos = pd.read_table(filename_or_file, delim_whitespace=True, names=["Ref", "Val", "Package", "PosX", "PosY", "Rot", "Side"], comment="#")
    pos.set_index("Ref", inplace=True)
    return pos

class KiCADFilenames(namedtuple("Filenames", ["pro", "pcb", "sch"])):
    """
    A class representing the top-level filenames of a KiCAD project.
    """
    def check_files_exist(self):
        """
        Check if all files in the Filenames tuple exist
        """
        for file in self:
            if not os.path.exists(file):
                raise FileNotFoundError(f"{file} does not exist.")
        return True

def kicad_project_filenames(proj_filename) -> KiCADFilenames:
    """
    Given a KiCAD project filename (or path),
    guess the PCB and schematic filenames.
    th, pos_outfile.name)
        pos_outfile.flush()
    """
    prefix, _ = os.path.splitext(proj_filename)
    return KiCADFilenames(
        prefix + ".kicad_pro",
        prefix + ".kicad_pcb",
        prefix + ".kicad_sch"
    )

def export_pos_from_pcb(pcb_filepath, pos_filepath, side=AssemblySide.Top, smd_only=False, exclude_through_hole=True, use_drill_file_origin=True):
    """
    Export position .pos file from a given .kicad_pcb file
    """
    extra_args = ""
    if smd_only:
        extra_args += " --smd-only"
    if exclude_through_hole:
        extra_args += " --exclude-fp-th"
    if side == AssemblySide.Top:
        extra_args += " --side front"
    elif side == AssemblySide.Bottom:
        extra_args += " --side back"
    if use_drill_file_origin:
        extra_args += " --use-drill-file-origin"

    cmd = f"kicad-cli pcb export pos '{pcb_filepath}' {extra_args} --units mm -o '{pos_filepath}'"
    run(cmd)

def export_bom_xml_from_sch(sch_filepath, xml_filepath):
    """
    Export position .pos file from a given .kicad_pcb file
    """
    run(f"kicad-cli sch export python-bom '{sch_filepath}' -o '{xml_filepath}'")

def export_and_read_position_file(pcb_filepath):
    """
    Export position .pos file from a given .kicad_pcb file
    """
    with tempfile.NamedTemporaryFile(mode='w+') as pos_outfile:
        export_pos_from_pcb(pcb_filepath, pos_outfile.name)
        pos_outfile.flush()
        # Read generated file
        pos_outfile.seek(0)
        pos_data = pos_outfile.read()
        return read_kicad_pos_file(StringIO(pos_data))

def export_and_read_bom_file(sch_filepath):
    """
    Export position .pos file from a given .kicad_pcb file
    """
    with tempfile.NamedTemporaryFile(mode='w+') as bom_outfile:
        export_bom_xml_from_sch(sch_filepath, bom_outfile.name)
        bom_outfile.flush()
        # Read generated file
        bom_outfile.seek(0)
        bom_data = bom_outfile.read()
        soup = BeautifulSoup(bom_data, 'xml')
        return soup

def text_or_None(elem):
    """
    Returns the text of an element if it exists, otherwise returns None.

    Args:
    elem: An element object.

    Returns:
    The text of the element if it exists, otherwise None.
    """
    if elem is None:
        return None
    else:
        return elem.text

def value_or_None(elem):
    """
    Returns the value attribute of an element if it exists, otherwise returns None.

    Args:
    elem: An element object.

    Returns:
    The value of the element if it exists, otherwise None.
    """
    if elem is None:
        return None
    else:
        return elem.get("value", None)

def component_list_to_dataframe(components):
    """
    Convert a list of Component objects to a Pandas DataFrame.
    The "Total count" column is automatically generated.
    """

    # Count how many times this specific part is used
    mpn_counter = Counter([c.mpn_or_value_plus_footprint for c in components if c.populate])

    # Create data frame from components
    df = pd.DataFrame([c.asdict() for c in components])

    # Build a map of RefDes => total count of this part
    mpn_total_count_by_refdes = {}
    for component in components:
        if not component.populate:
            continue
        mpn = component.mpn_or_value_plus_footprint
        mpn_total_count_by_refdes[component.refdes] = mpn_counter.get(mpn, 0)

    # Add total MPN counter column
    df['Total count'] = df['Ref'].map(lambda ref: mpn_total_count_by_refdes.get(ref, "NaN"))
    return df


import re
from collections import namedtuple
from dataclasses import dataclass
from os import linesep

blueprint_pattern = re.compile(
    r"Blueprint (\d+): "
    + r"Each ore robot costs (\d+) ore. "
    + r"Each clay robot costs (\d+) ore. "
    + r"Each obsidian robot costs (\d+) ore and (\d+) clay. "
    + r"Each geode robot costs (\d+) ore and (\d+) obsidian."
)

@dataclass(frozen=True)
class Blueprint:
    id: int
    ores_orebot: int
    ores_claybot: int
    ores_obsidianbot: int
    clays_obsidianbot: int
    ores_geodebot: int
    obsidians_geodebot: int

def parse(input: str):
    matches = blueprint_pattern.findall(input)

    blueprints = [Blueprint(*map(int, match)) for match in matches]

    return blueprints

def part_1(blueprints: list[Blueprint]):
    for b in blueprints:
        states = []

        states.append((24, 1, 0, 0, 0, 0, 0, 0, 0))

        most_ore_exp_bot = max(
            b.ores_orebot, b.ores_claybot, b.ores_obsidianbot, b.ores_geodebot
        )

        i = 0
        while states:
            state = states.pop(0)
            (
                minutes_left,
                ore_bots,
                clay_bots,
                obsidian_bots,
                geode_bots,
                ores,
                clays,
                obsidians,
                geodes,
            ) = state

            i += 1
            if i % 1_000 == 0:
                print(len(states), state)

            new_minutes_left = minutes_left - 1
            new_ores = ores + ore_bots
            new_clays = clays + clay_bots
            new_obsidians = obsidians + obsidian_bots
            new_geodes = geodes + geode_bots

            def push_state(
                new_minutes_left=new_minutes_left,
                ore_bots=ore_bots,
                clay_bots=clay_bots,
                obsidian_bots=obsidian_bots,
                geode_bots=geode_bots,
                ores=new_ores,
                clays=new_clays,
                obsidians=new_obsidians,
                geodes=new_geodes,
            ):
                states.append(
                    (
                        new_minutes_left,
                        ore_bots,
                        clay_bots,
                        obsidian_bots,
                        geode_bots,
                        ores,
                        clays,
                        obsidians,
                        geodes,
                    )
                )

            bought_ore, bought_clay, bought_obs, bought_geode = (
                False,
                False,
                False,
                False,
            )

            if ore_bots < most_ore_exp_bot and ores >= b.ores_orebot:
                bought_ore = True
                push_state(
                    ore_bots=ore_bots + 1,
                    ores=new_ores - b.ores_orebot,
                )

            if ores >= b.ores_claybot:
                bought_clay = True
                push_state(
                    clay_bots=clay_bots + 1,
                    ores=new_ores - b.ores_claybot,
                )

            if ores >= b.ores_obsidianbot and clays >= b.clays_obsidianbot:
                bought_obs = True
                push_state(
                    obsidian_bots=obsidian_bots + 1,
                    ores=new_ores - b.ores_obsidianbot,
                    clays=new_clays - b.clays_obsidianbot,
                )

            if ores >= b.ores_geodebot and obsidians >= b.obsidians_geodebot:
                bought_geode = True
                push_state(
                    obsidian_bots=obsidian_bots + 1,
                    ores=new_ores - b.ores_geodebot,
                    obsidians=new_obsidians - b.obsidians_geodebot,
                )

            # do nothing
            if not (bought_ore and bought_clay and bought_obs and bought_geode):
                push_state()


import logging

import boto3

from app import settings
from app.exceptions.upload_file import S3SendFileException

class S3Service:
    def __init__(self, bucket_name, aws_access_key_id, aws_secret_access_key):
        self._bucket_name = bucket_name
        self._aws_access_key_id = aws_access_key_id
        self._aws_secret_access_key = aws_secret_access_key

    def upload_file(self, content: bytes, file_name: str) -> None:

        try:
            s3 = boto3.client(
                's3', aws_access_key_id=self._aws_access_key_id, aws_secret_access_key=self._aws_secret_access_key
            )
            s3.upload_fileobj(content, self._bucket_name, file_name)
        except Exception as e:
            logging.error(f'Error send file to s3 {e}')
            raise S3SendFileException(e)

def get_s3_service() -> S3Service:
    return S3Service(
        bucket_name=settings.AWS_S3_BUCKET_NAME,
        aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
        aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
    )


import cv2
from matplotlib import pyplot as plt

# Download files from https://drive.google.com/file/d/1XdZLvORnCnfpyBYflh15I58VQrQdVlUe/view?usp=sharing

h_skin_hist = 0
h_nonskin_hist = 0
s_skin_hist = 0
s_nonskin_hist = 0
for im_id in range(1, 4):
    print(im_id)
    im = cv2.imread("SkinDetection\SkinTrain" + str(im_id) + ".jpg")
    mask = cv2.imread("SkinDetection\SkinTrain" + str(im_id) + "_mask.jpg", 0)
    im_hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
    h_skin_hist = h_skin_hist + cv2.calcHist([im_hsv], [0], mask, [256], [0, 256])
    h_nonskin_hist = h_nonskin_hist + cv2.calcHist([im_hsv], [0], 255 - mask, [256], [0, 256])
    s_skin_hist = s_skin_hist + cv2.calcHist([im_hsv], [1], mask, [256], [0, 256])
    s_nonskin_hist = s_nonskin_hist + cv2.calcHist([im_hsv], [1], 255 - mask, [256], [0, 256])

h_skin_prob = h_skin_hist / sum(h_skin_hist) # P(h|skin)
h_nonskin_prob = h_nonskin_hist / sum(h_nonskin_hist) # P(h|non-skin)
s_skin_prob = s_skin_hist / sum(s_skin_hist) # P(s|skin)
s_nonskin_prob = s_nonskin_hist / sum(s_nonskin_hist) # P(s|non-skin)

plt.plot(h_skin_prob, 'r', label = 'P(h|skin)')
plt.plot(h_nonskin_prob, 'b', label = 'P(h|non-skin)')
plt.legend()
plt.figure()
plt.plot(s_skin_prob, 'r', label = 'P(s|skin)')
plt.plot(s_nonskin_prob, 'b', label = 'P(s|non-skin)')
plt.show()


import torch

from torch.nn import Module, Embedding, Linear, MultiheadAttention, LayerNorm, Dropout
from pykt.models.utils import transformer_FFN, pos_encode, ut_mask, get_clones
from torch.nn.functional import binary_cross_entropy

class SAKT(Module):
    def __init__(self, num_c, num_q, seq_len, emb_size, num_attn_heads, dropout, num_en=2, emb_type="qid"):
        super().__init__()
        self.model_name = "sakt"
        self.emb_type = emb_type

        self.num_c = num_c
        self.num_q = num_q
        self.seq_len = seq_len
        self.emb_size = emb_size
        self.num_attn_heads = num_attn_heads
        self.dropout = dropout
        self.num_en = num_en

        # num_c, seq_len, emb_size, num_attn_heads, dropout, emb_path="")
        self.q_emb = Embedding(num_q, emb_size)
        self.c_emb = Embedding(num_c+1, emb_size, padding_idx=0)
        # self.P = Parameter(torch.Tensor(self.seq_len, self.emb_size))
        self.position_emb = Embedding(seq_len, 4*emb_size)

        self.blocks = get_clones(Blocks(emb_size, num_attn_heads, dropout), self.num_en)

        self.dropout_layer = Dropout(dropout)
        self.pred = Linear(self.emb_size, 1)
        self.x_emb = Linear(2*emb_size, emb_size)
        self.e_emb = Linear(4*emb_size, emb_size)

    def multi_skills_embedding(self,c):
        concept_emb_sum = self.c_emb(c + 1).sum(-2)
        # [batch_size, seq_len,1]
        concept_num = torch.where(c < 0, 0, 1).sum(-1).unsqueeze(-1)
        concept_num = torch.where(concept_num == 0, 1, concept_num)
        concept_avg = (concept_emb_sum / concept_num)
        return concept_avg

    def base_emb(self, q, c, r, qry, cry):
        q_emb = self.q_emb(q)
        c_emb = self.multi_skills_embedding(c)
        x = torch.cat([q_emb, c_emb], -1)
        r = r.unsqueeze(-1)
        resp_t = r.repeat(1, 1, 2 * self.emb_size)
        resp_f = (1 - r).repeat(1, 1, 2 * self.emb_size)
        e_emb = torch.cat([x.mul(resp_t), x.mul(resp_f)], -1)

        qshftemb = torch.cat([self.q_emb(qry), self.multi_skills_embedding(cry)], -1)

        posemb = self.position_emb(pos_encode(e_emb.shape[1]))
        e_emb = e_emb + posemb
        return self.x_emb(qshftemb), self.e_emb(e_emb)

    def forward(self, data):
        q, c, r, t = data["qseqs"].long(), data["cseqs"].long(), data["rseqs"].long(), data["tseqs"].long()
        qshft, cshft, rshft, tshft = data["shft_qseqs"], data["shft_cseqs"], data["shft_rseqs"], data["shft_tseqs"]
        m, sm = data["masks"], data["smasks"]

        qshftemb, xemb = self.base_emb(q, c, r, qshft, cshft)

        # print(f"qemb: {qemb.shape}, xemb: {xemb.shape}, qshftemb: {qshftemb.shape}")
        for i in range(self.num_en):
            xemb = self.blocks[i](qshftemb, xemb, xemb)
        p = torch.sigmoid(self.pred(self.dropout_layer(xemb))).squeeze(-1)

        y = torch.masked_select(p, sm)
        q = torch.masked_select(qshft, sm)
        t = torch.masked_select(rshft, sm)
        loss = binary_cross_entropy(y.double(), t.double())
        return loss, y, t, q + self.num_q * t

class Blocks(Module):
    def __init__(self, emb_size, num_attn_heads, dropout) -> None:
        super().__init__()

        self.attn = MultiheadAttention(emb_size, num_attn_heads, dropout=dropout)
        self.attn_dropout = Dropout(dropout)
        self.attn_layer_norm = LayerNorm(emb_size)

        self.FFN = transformer_FFN(emb_size, dropout)
        self.FFN_dropout = Dropout(dropout)
        self.FFN_layer_norm = LayerNorm(emb_size)

    def forward(self, q=None, k=None, v=None):
        q, k, v = q.permute(1, 0, 2), k.permute(1, 0, 2), v.permute(1, 0, 2)
        # attn -> drop -> skip -> norm
        # transformer: attn -> drop -> skip -> norm transformer default
        causal_mask = ut_mask(seq_len=k.shape[0])
        attn_emb, _ = self.attn(q, k, v, attn_mask=causal_mask)

        attn_emb = self.attn_dropout(attn_emb)
        attn_emb, q = attn_emb.permute(1, 0, 2), q.permute(1, 0, 2)

        attn_emb = self.attn_layer_norm(q + attn_emb)

        emb = self.FFN(attn_emb)
        emb = self.FFN_dropout(emb)
        emb = self.FFN_layer_norm(attn_emb + emb)
        return emb