import os
import numpy as np
import time

import pybullet as p

from logger import logger

from robot.conveyor import ConveyorBelt
from robot.pallet import Pallet
from packer.packer import PalletPacker


class UrdfManager:
    urdf_filename = None
    urdf_filename_tool = "welding_gun"
    physics_client = None
    body_id = None

    def __init__(self, robot_start_position, application_path, palletizing):
        self.robot_start_position = robot_start_position
        self.application_path = application_path
        self.palletizing = palletizing

        self.conveyor = None
        self.pallet = None

    def start_loop(self, urdf):
        self.urdf_filename = urdf

        p.resetSimulation()
        self.load_models()
        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
        self.state_id = p.saveState()

    def run_pybullet(self, type="DIRECT"):
        self.physics_client = p.connect(getattr(p, type))
        logger.info(f"Connect to {self.physics_client} by {type}")

    def load_models(self):
        p.loadURDF(
            os.path.join(self.application_path, "urdf", "plane.urdf"),
            physicsClientId=self.physics_client,
        )

        urdf_path = os.path.join(
            self.application_path, "urdf", f"{self.urdf_filename}.urdf"
        )
        self.body_id = p.loadURDF(
            urdf_path,
            self.robot_start_position,
            useFixedBase=True,
        )

    def load_palletizing(self):
        if self.conveyor:
            self.conveyor = None
        if self.pallet:
            self.pallet = None
        self.start_loop(self.urdf_filename)
        
        k = 0.2
        self.pallet_target = [1.3 + k, 0.6, 0.1]
        self.pallet_size = [0.3 + k, 1.2, 0.144]
        self.pallet_position = [
            (self.pallet_target[i] - self.pallet_size[i] * 0.5)
            for i in range(len(self.pallet_target))
        ]
        self.pallet = Pallet(self.pallet_position)

        self.box_size = [0.2, 0.2, 0.2]

        # положение дальнего угла
        self.conv_target = [1.1 + 0.5, -0.63, 0.4]
        self.conv_size = [0.5, 3.0, 0.1]
        self.conv_direction = self.conv_size.index(max(self.conv_size))
        self.conv_pos = [
            (self.conv_target[i] - self.conv_size[i] * 0.5)
            for i in range(len(self.conv_target))
        ]
        self.conveyor = ConveyorBelt(
            self.conv_size,
            self.conv_pos,
            self.conv_direction,
            self.box_size,
            0.2,
            0.1,
            self.palletizing,
        )

        self.packing = PalletPacker(
            {
                "x": self.pallet_size[0],
                "y": self.pallet_size[1],
                "z": self.box_size[2],  # высота паллеты в один слой
            },
            {
                "x": self.box_size[0],
                "y": self.box_size[1],
                "z": self.box_size[2],
            },
        ).pack()
        self.packing = sorted(self.packing, key=lambda x: (x[0], x[1]))

        for p in self.packing:
            # logger.info(p)
            
            for i in range(len(self.box_size)):
                p[i] = (
                    self.pallet_target[i]
                    - p[i]
                    - self.box_size[i] * 0.5
                )
            dir_wall = (self.conv_direction + 1) % (3 - 1)
            p[dir_wall] -= self.box_size[dir_wall]
            p[self.conv_direction] -= self.box_size[self.conv_direction] * 0.5
            p[2] += self.pallet_size[2]
        # logger.info(f"{self.packing[0]} {self.pallet_target}")

        time.sleep(1)

    def get_pybullet_image(self):
        if self.physics_client is None:
            return

        width, height, rgb, _, _ = p.getCameraImage(
            width=500,
            height=500,
            viewMatrix=p.computeViewMatrix(
                cameraEyePosition=[4, -4, 1.5],
                # cameraEyePosition=[4,-4,4],
                cameraTargetPosition=[0, 0, 0],  # Центр фокусировки камеры
                cameraUpVector=[0, 0, 1],  # Направление вверх
            ),
            projectionMatrix=p.computeProjectionMatrixFOV(
                fov=60.0, aspect=1.0, nearVal=0.1, farVal=10.0
            ),
            # renderer=p.ER_TINY_RENDERER,
            physicsClientId=self.physics_client,
        )

        return (rgb, width, height)