import torch
import torch.nn as nn 
import math 
from lift3d.models.concept.MetaWorld.utils import rotate
from lift3d.models.concept.MetaWorld.utils import Para_Estimator, Knowledge_Encoder, PT_Knowledge_Encoder, PN_Knowledge_Encoder, Fusion_Knowledge_Encoder
from lift3d.models.concept.MetaWorld.Dense_Concept import Dense_Cuboid, Dense_Handle

def Cuboid(size, position, rotation, Sema):
    """
    :size[B,3]: [0]=length, [1]=width, [2]=height
    :position[B,3], rotation[B,3]:
    """
    B = size.shape[0]
    
    vertices = torch.tensor([
        [-1 / 2, 1 / 2, 1 / 2],
        [1 / 2, 1 / 2, 1 / 2],
        [-1 / 2, 1 / 2, -1 / 2],
        [1 / 2, 1 / 2, -1 / 2],
        [-1 / 2, -1 / 2, 1 / 2],
        [1 / 2, -1 / 2, 1 / 2],
        [-1 / 2, -1 / 2, -1 / 2],
        [1 / 2, -1 / 2, -1 / 2]
    ], device=position.device, dtype=position.dtype)
    
    deformation = size.unsqueeze(dim=1).repeat(1, 8, 1)
    
    vertices = rotate(vertices * deformation, rotation) + position.unsqueeze(dim=1)
    Semantic = torch.ones([B, 8, 1], device=position.device, dtype=position.dtype) * Sema
    Affordance = torch.cat((vertices, Semantic), dim=-1)
    
    return Affordance

def Handle(inner, outer, height, position, rotation, Sema):
    """
    :param inner[B,2]: [0]=inner_length, [1]=inner_width
    :param outer[B,2]: [0]=outer_length, [1]outer_width
    :height[B,1]: height of Handle
    :position[B,3], rotation[B,3]: position and rotation
    """
    B = inner.shape[0]

    # [B, 16, 3] vertex template
    vertices = torch.tensor([
        [-0.5,  0.5,  0.5], [0.5,  0.5,  0.5],
        [-0.5, -0.5,  0.5], [0.5, -0.5,  0.5],
        [-0.5,  0.5, -0.5], [0.5,  0.5, -0.5],
        [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5],
        [-0.5,  0.5,  0.5], [0.5,  0.5,  0.5],
        [-0.5,  0.5, -0.5], [0.5,  0.5, -0.5],
        [-0.5, -0.5,  0.5], [0.5, -0.5,  0.5],
        [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5]
    ], device=inner.device, dtype=inner.dtype).unsqueeze(0).repeat(B, 1, 1)  # [B, 16, 3]

    # [B, 1, 3] → [B, 8, 3] for outer & inner box
    outer_box1 = torch.cat([outer[:, 0:1], outer[:, 1:2], height], dim=-1).unsqueeze(1).repeat(1, 4, 1)
    outer_box2 = torch.cat([outer[:, 0:1], inner[:, 1:2], height], dim=-1).unsqueeze(1).repeat(1, 4, 1)
    inner_box = torch.cat([inner[:, 0:1], inner[:, 1:2], height], dim=-1).unsqueeze(1).repeat(1, 8, 1)
    deformation = torch.cat([outer_box1, outer_box2, inner_box], dim=1)  # [B, 16, 3]

    # Apply scale, rotation, and translation
    vertices = rotate(vertices * deformation, rotation) + position.unsqueeze(1)  # [B, 16, 3]
    Semantic = torch.ones([B, 16, 1], device=inner.device, dtype=inner.dtype) * Sema
    Affordance = torch.cat((vertices, Semantic), dim=-1)

    return Affordance

def Concept_DrawerOpen (paras):
    
    B = paras.shape[0]
    zeros = torch.zeros([B,1], device=paras.device, dtype=paras.dtype)
    ones = torch.ones([B,1], device=paras.device, dtype=paras.dtype)
    
    # Drawer Box
    size_Box = paras[:, 0:3]
    position_Box = paras[:, 3:6]
    rotation = paras[:, 6:9]
    vertices_Box = Cuboid(size_Box, position_Box, rotation, 1)
    
    # Deawer Handle
    outer = paras[:, 9:11]
    inner = outer - paras[:, 11:13]
    height = paras[:, 13:14]
    position_Handle = position_Box + torch.cat([0.5 * size_Box[:, 0:1] + 0.5 * inner[:, 0:1], zeros, zeros], dim=-1)
    vertices_Handle = Handle(inner, outer, height, position_Handle, rotation, 2)
    
    direction = rotate(torch.cat([-1 * ones, zeros, zeros], dim=-1).unsqueeze(dim=1), rotation).view([B, 3])
    Grasp_Pose = torch.cat([position_Handle, direction], dim=-1)
    
    Affordance = torch.cat([vertices_Box, vertices_Handle], dim=1)
    return Affordance, Grasp_Pose

def process_paras(raw):
    """
    Normalize raw output from Para_Estimator into structured and bounded physical parameters.
    Input: raw tensor of shape [B, 14], unrestricted range
    Output: tensor of shape [B, 14], scaled to valid physical ranges
    """
    B = raw.shape[0]

    # Drawer Box size: [0.1, 0.5] (x,y,z)
    size_box = torch.sigmoid(raw[:, 0:3]) * 0.4 + 0.1

    # Drawer Box position: [-2, 2]
    pos_box = torch.tanh(raw[:, 3:6]) * 2.0

    # Rotation: [-pi, pi]
    rotation = torch.tanh(raw[:, 6:9]) * math.pi

    # Handle outer size (w, h): [0.02, 0.15]
    outer = torch.sigmoid(raw[:, 9:11]) * 0.13 + 0.02

    # Handle thickness (t_w, t_h): [0.005, 0.03]
    thickness = torch.sigmoid(raw[:, 11:13]) * 0.025 + 0.005
    inner = outer - 2 * thickness  # Make sure inner < outer

    # Handle height (depth along drawer axis): [0.001, 0.01]
    height = torch.sigmoid(raw[:, 13:14]) * 0.009 + 0.001

    return torch.cat([
        size_box,
        pos_box,
        rotation,
        outer,
        thickness,
        height
    ], dim=-1)

class Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = Knowledge_Encoder(Knowledge_feature_dim, Knowledge_attension_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        paras = process_paras(paras)
        knowledge, grasp_pose = Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class PN_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(PN_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = PN_Knowledge_Encoder(Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        # paras = process_paras(paras)
        knowledge, grasp_pose = Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class Fusion_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_Semantic_dim = paralist[2]
        Knowledge_xyz_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(Fusion_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = Fusion_Knowledge_Encoder(Knowledge_Semantic_dim, 16+8, Knowledge_xyz_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        paras = process_paras(paras)
        knowledge, grasp_pose = Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class PT_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(PT_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = PT_Knowledge_Encoder(Knowledge_feature_dim, Knowledge_attension_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        paras = process_paras(paras)
        knowledge, grasp_pose = Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
#######################################################################

# Dense point cloud in concept space 

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

def Dense_Concept_DrawerOpen (paras):
    
    B = paras.shape[0]
    zeros = torch.zeros([B,1], device=paras.device, dtype=paras.dtype)
    ones = torch.ones([B,1], device=paras.device, dtype=paras.dtype)
    
    # Drawer Box
    size_Box = paras[:, 0:3]
    position_Box = paras[:, 3:6]
    rotation = paras[:, 6:9]
    vertices_Box = Dense_Cuboid(size_Box, position_Box, rotation, Sem=1, N=40)
    
    # Deawer Handle
    outer = paras[:, 9:11]
    inner = outer - paras[:, 11:13]
    height = paras[:, 13:14]
    position_Handle = position_Box + torch.cat([0.5 * size_Box[:, 0:1] + 0.5 * inner[:, 0:1], zeros, zeros], dim=-1)
    vertices_Handle = Dense_Handle(inner, outer, height, position_Handle, rotation, Sem=2, order=2, N=80)
    
    direction = rotate(torch.cat([-1 * ones, zeros, zeros], dim=-1).unsqueeze(dim=1), rotation).view([B, 3])
    Grasp_Pose = torch.cat([position_Handle, direction], dim=-1)
    
    Affordance = torch.cat([vertices_Box, vertices_Handle], dim=1)
    return Affordance, Grasp_Pose

class Dense_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(Dense_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = Knowledge_Encoder(Knowledge_feature_dim, Knowledge_attension_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        # paras = process_paras(paras)
        knowledge, grasp_pose = Dense_Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class PN_Dense_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(PN_Dense_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = PN_Knowledge_Encoder(Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        # paras = process_paras(paras)
        knowledge, grasp_pose = Dense_Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class Fusion_Dense_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_Semantic_dim = paralist[2]
        Knowledge_xyz_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(Fusion_Dense_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = Fusion_Knowledge_Encoder(Knowledge_Semantic_dim, 40+80, Knowledge_xyz_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        # paras = process_paras(paras)
        knowledge, grasp_pose = Dense_Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state
    
class PT_Dense_Concept_Module_DrawerOpen(nn.Module):
    def __init__ (self, in_dim, paralist):
        estimator_hidden_dim = paralist[0]
        Geo_para_dim = paralist[1]
        Knowledge_feature_dim = paralist[2]
        Knowledge_attension_dim = paralist[3]
        Concept_out_dim = paralist[4]
        super(PT_Dense_Concept_Module_DrawerOpen, self).__init__()
        self.para_estimator = Para_Estimator(in_dim, estimator_hidden_dim, Geo_para_dim)
        self.knowledge_encoder = PT_Knowledge_Encoder(Knowledge_feature_dim, Knowledge_attension_dim, Concept_out_dim)
        
    def forward(self, feature, robot_state):
        paras = self.para_estimator(feature)
        # paras = process_paras(paras)
        knowledge, grasp_pose = Dense_Concept_DrawerOpen(paras)
        concept_emb = self.knowledge_encoder(knowledge)
        robot_state = torch.cat((grasp_pose, robot_state), dim=-1)
        return concept_emb, robot_state