from .llava_vla_model import *
import argparse
import numpy as np
from PIL import Image

EXPECTED_ACTION_DIM = 70
NUM_ACTIONS = 5
ACTION_DIM = 14

import torch
import torchvision.transforms as T
from PIL import Image
import os


def encode_obs(observation):
    input_rgb_arr = [        
        observation["observation"]["left_camera"]["rgb"],
        observation["observation"]["right_camera"]["rgb"],
        observation["observation"]["head_camera"]["rgb"],
        observation["observation"]["front_camera"]["rgb"],
    ]
    input_state = observation["joint_action"]["vector"]

    return input_rgb_arr, input_state


def combined_images(images):

    images = [np.array(img) if not isinstance(img, np.ndarray) else img for img in images]

    height, width, channels = images[0].shape

    if channels != 3:
        raise ValueError(f"Expected 3 channels (RGB/BGR), got {channels}")


    images = [(img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8) for img in images]
    images = [img[:, :, ::-1] if img[0, 0, 2] > img[0, 0, 0] else img for img in images]  
    stitched = np.zeros((2 * height, 2 * width, 3), dtype=np.uint8)

    # images[0] (left_camera)
    stitched[0:height, 0:width] = images[0]
    #  images[1] (right_camera)
    stitched[0:height, width:2*width] = images[1]
    #  images[2] (head_camera)
    stitched[height:2*height, 0:width] = images[2]
    #  images[3] (front_camera)
    stitched[height:2*height, width:2*width] = images[3]

    # RGB 
    combine_image = Image.fromarray(stitched)

    return combine_image

def get_model(usr_args):
    if isinstance(usr_args, dict):
        usr_args = argparse.Namespace(**usr_args)
    return LLaVA(usr_args)

def eval(TASK_ENV, model, observation):
    
    instruction = TASK_ENV.get_instruction()
    model.set_language(instruction)
    input_rgb_arr, input_state = encode_obs(observation)
  
    image=combined_images(input_rgb_arr)
    input_ids,images=model.compose_robot_input(image,instruction,input_state)
    actions = model.get_action(input_ids, images)
    
    print("Action shape before reshape:", actions.shape)
    
    if actions.shape == (70,):
        actions = actions.reshape(5, 14)  

    for i in range(actions.shape[0]):
        action = actions[i]  
        TASK_ENV.take_action(action)  
        observation = TASK_ENV.get_obs()
        input_rgb_arr, input_state = encode_obs(observation)
        print(f"input_state_after_action (step {i}):", input_state)
    
    return observation


def reset_model(model):  
    # Clean the model cache at the beginning of every evaluation episode, such as the observation window
    pass