import os
import torch
from roboimi.utils.utils import load_data, set_seed
from roboimi.detr.policy import ACTPolicy, CNNMLPPolicy, ACTTVPolicy
from roboimi.gr00t.policy import gr00tPolicy

class ModelInterface:
    def __init__(self, config):
        self.config = config

    def setup(self):
        set_seed(self.config['seed'])
        
        self.setup_task_parameters()
        self.setup_policy_parameters()

    def setup_task_parameters(self):
        is_sim = self.config['task_name'][:4] == 'sim_'   #attention here
        if is_sim:
            from roboimi.utils.constants import SIM_TASK_CONFIGS
            task_config = SIM_TASK_CONFIGS[self.config['task_name']]
        
        self.config['dataset_dir'] = task_config['dataset_dir']
        self.config['num_episodes'] = task_config['num_episodes']
        self.config['episode_len'] = task_config['episode_len']
        self.config['camera_names'] = task_config['camera_names']

    def setup_policy_parameters(self):
        if self.config['policy_class'] == 'ACT':
            self.config['policy_config'] = {
                'lr': self.config['lr'],
                'num_queries': self.config['chunk_size'],
                'kl_weight': self.config['kl_weight'],
                'hidden_dim': self.config['hidden_dim'],
                'dim_feedforward': self.config['dim_feedforward'],
                'lr_backbone': self.config['lr_backbone'],
                'backbone': self.config['backbone'],
                'enc_layers': self.config['enc_layers'],
                'dec_layers': self.config['dec_layers'],
                'nheads': self.config['nheads'],
                'camera_names': self.config['camera_names'],
            }
        elif self.config['policy_class'] == 'ACTTV':
            backbone = 'dino_v2'
            self.config['policy_config'] = {
                'lr': self.config['lr'],
                'num_queries': self.config['chunk_size'],
                'kl_weight': self.config['kl_weight'],
                'hidden_dim': self.config['hidden_dim'],
                'dim_feedforward': self.config['dim_feedforward'],
                'lr_backbone': self.config['lr_backbone'],
                'backbone': backbone,
                'enc_layers': self.config['enc_layers'],
                'dec_layers': self.config['dec_layers'],
                'nheads': self.config['nheads'],
                'camera_names': self.config['camera_names'],
                'state_dim': 16,
                'action_dim': 16,
                'qpos_noise_std': self.config['qpos_noise_std'],
            }
        elif self.config['policy_class'] == 'CNNMLP':
            self.config['policy_config'] = {
                'lr': self.config['lr'],
                'lr_backbone': self.config['lr_backbone'],
                'backbone': self.config['backbone'],
                'num_queries': 1,
                'camera_names': self.config['camera_names'],
            }
        elif self.config['policy_class'] == 'GR00T':
            # GR00T uses its own config section from config.yaml
            gr00t_config = self.config.get('gr00t', {})
            self.config['policy_config'] = {
                'lr': gr00t_config.get('lr', self.config['lr']),
                'lr_backbone': gr00t_config.get('lr_backbone', self.config['lr_backbone']),
                'weight_decay': gr00t_config.get('weight_decay', 1e-4),
                'embed_dim': gr00t_config.get('embed_dim', 1536),
                'hidden_dim': gr00t_config.get('hidden_dim', 1024),
                'state_dim': gr00t_config.get('state_dim', 16),
                'action_dim': gr00t_config.get('action_dim', 16),
                'num_queries': gr00t_config.get('num_queries', 16),
                'num_layers': gr00t_config.get('num_layers', 16),
                'nheads': gr00t_config.get('nheads', 32),
                'mlp_ratio': gr00t_config.get('mlp_ratio', 4),
                'dropout': gr00t_config.get('dropout', 0.2),
                'backbone': gr00t_config.get('backbone', 'dino_v2'),
                'position_embedding': gr00t_config.get('position_embedding', 'sine'),
                'camera_names': self.config['camera_names'],
            }
        else:
            raise NotImplementedError

    def make_policy(self):
        if self.config['policy_class'] == 'ACT':
            return ACTPolicy(self.config['policy_config'])
        elif self.config['policy_class'] == 'ACTTV':
            return ACTTVPolicy(self.config['policy_config'])
        elif self.config['policy_class'] == 'CNNMLP':
            return CNNMLPPolicy(self.config['policy_config'])
        elif self.config['policy_class'] == 'GR00T':
            return gr00tPolicy(self.config['policy_config'])
        else:
            raise NotImplementedError

    def make_optimizer(self, policy):
        return policy.configure_optimizers()

    def load_data(self):
        return load_data(
            self.config['dataset_dir'], 
            self.config['num_episodes'], 
            self.config['camera_names'], 
            self.config['batch_size'], 
            self.config['batch_size']
        )

    def load_policy(self, ckpt_name):
        policy = self.make_policy()
        ckpt_path = os.path.join(self.config['ckpt_dir'], ckpt_name)
        policy.load_state_dict(torch.load(ckpt_path))
        policy.cuda()
        policy.eval()
        return policy

    # def make_env(self):
    #     if self.config['real_robot']:
    #         from aloha_scripts.real_env import make_real_env
    #         return make_real_env(init_node=True)
    #     else:
    #         from sim_env import make_sim_env
    #         return make_sim_env(self.config['task_name'])