import gc
import pickle

import os
import time

os.environ["TOKENIZERS_PARALLELISM"] = "false"

os.environ['DEVICE'] = "cuda"
os.environ["WANDB_DISABLED"] = "true"

from data_utils.dataset import load_data  # data functions
from data_utils.dataset import compute_dict_mean, set_seed  # helper functions
from policy_heads import *
# from data_utils.lerobot_dataset import load_data
from aloha_scripts.constants import TASK_CONFIGS
from dex_vla.utils.robot_data_processor import DexVLAProcess
# from paligemma_vla.utils.robot_data_processor import PaliGemmaVLAProcess
from transformers import AutoConfig, AutoModel, AutoProcessor
from dex_vla import DexVLATrainer
from data_utils.data_collator import *

import IPython
e = IPython.embed
from data_utils.data_collator import DexVLADataCollatorForSupervisedDataset, PaliGemmaVLADataCollatorForSupervisedDataset
from dex_vla import model_load_utils as ml_utils
import torch
local_rank = None
from aloha_scripts.utils import *
#  >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>parameters<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
@dataclass
class ActionHeadArguments:
    policy_head_type: str = field(default="dit_diffusion_policy") # unet_diffusion_policy
    policy_head_size: str = field(default="DiT_B") # DiT_L, DiT_XL, DiT_B, DiT_S
    state_dim: int = 7
    action_dim: int = 10


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
    version: Optional[str] = field(default="v0")
    model_pretrain: Optional[str] = field(default="")  # pretrained model weights path
    from_scratch: bool = field(default=False)

    external_vision_encoder: Optional[str] = field(default="None")

    concat: str = field(default="None")
    policy_class: str = field(default="droid_diffusion")

    # with_external_vit: bool = field(default=False)
    with_llm_head: bool = field(default=False)
    with_text_fcs: bool = field(default=False)
    only_using_input_embeddings: bool = field(default=False)  # using only input embeddings
    using_film: bool = field(default=False) # fusion modules
    using_xattn: bool = field(default=False) # fusion modules

    using_state: bool = field(default=False) # input states into VLM

    using_channel_cat: bool = field(default=False)
    using_all_reasoning_hidden: bool = field(default=False)
    ground_truth_reasoning: bool = field(default=False)

    Using_EMA_Pretrain_DiT: bool = field(default=False)

    load_pretrain_dit: bool = field(default=False) # loading pretrained dit weights
    pretrain_dit_path: Optional[str] = field(default=None) # path to pretrained dit weights

    freeze_policy_head: bool = field(default=False)
    is_tinyvla: bool = field(default=False)
    using_joint_attn: bool = field(default=False)

    # vla_model_type: Optional[str] = field(default='dex_vla')

@dataclass
class DataArguments:
    # model_name_or_path: Optional[str] = field(default="facebook/opt-125m") # equals to base model path when set load_pretrain=True
    # model_pretrain: Optional[str] = field(default="") # pretrained model weights path
    lazy_preprocess: bool = False
    episode_first: bool = True  # batchsampler will samples episode index first and then samples timesteps
    select_seg_token_mask: bool = False
    use_reasoning: bool = False
    is_multimodal: bool = False
    image_aspect_ratio: str = 'square'
    task_name: str = field(default="stack_cube_2024_6_2")
    skip_mirrored_data: bool = field(default=False)
    chunk_size: int = field(default=16)
    delta_control: bool = field(default=False)
    image_size_stable: str = "480"  # default 270 x 480 and pretrain may be 180 x 320
    image_size_wrist: str = "56" # specify the image size of wrist camera
    history_images_length: int = 1
    home_lerobot: str = '/media/rl/HDD/data/data/aloha_data/lerobot'

@dataclass
class TrainingArguments(transformers.TrainingArguments):
    using_ema: bool = field(default=False) # whether to use ema update whole module

    local_debug: bool = field(default=False)

    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    adam_beta1: float = field(default=0.9)
    adam_beta2: float = field(default=0.98)
    adam_epsilon: float = field(default=1e-7)
    remove_unused_columns: bool = field(default=False)

    flash_attn: bool = field(default=False)

    freeze_vision_tower: bool = field(default=False)
    freeze_backbone: bool = field(default=False)
    tune_mm_mlp_adapter: bool = field(default=False)
    resume_from_checkpoint: bool = field(default=False)
    llm_loss_weight: float = field(default=1.0)

    seed: int = field(default=0)

    # logger
    logging_dir: str = field(default='./logs')  # TensorBoard日志的保存目录
    logging_strategy: str = field(default='steps')  # 设置为`steps`表示每几步记录一次日志
    logging_steps: int = field(default=10)

    save_steps: int = field(default=10)  # 每隔多少步保存一次模型
    num_train_epochs: int = field(default=3)
    max_steps: int = field(default=5000)

    # validate
    do_eval: bool = field(default=False)
    evaluation_strategy: str = field(default="no")
    eval_steps: int = field(default=200)
    per_device_eval_batch_size: int = field(default=32)

    load_pretrain: bool = False

    dataloader_pin_memory: bool = False
    # lora
    lora_enable: bool = False
    lora_module: str = "vit"
    lora_task_type: str = 'CAUSAL_LM'
    lora_r: int = 64
    lora_alpha: int = 256
    lora_dropout: float = 0.05
    lora_weight_path: str = ""
    lora_bias: str = "none"
    non_lora_lr: Optional[float] = None
    group_by_modality_length: bool = field(default=False)

    model_max_length: int = field(
        default=2048,
        metadata={
            "help":
                "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    double_quant: bool = field(
        default=True,
        metadata={"help": "Compress the quantization statistics through double quantization."}
    )
    quant_type: str = field(
        default="nf4",
        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
    )
    bits: int = field(
        default=16,
        metadata={"help": "How many bits to use."}
    )


#  <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<parameters>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

def rank0_print(*args):
    if local_rank == 0:
        print(*args)

def parse_param():
    global local_rank

    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments, ActionHeadArguments))
    model_args, data_args, training_args, action_head_args = parser.parse_args_into_dataclasses()

    local_rank = training_args.local_rank
    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))

    #     print("##"*50)
    #     print(training_args.logging_dir)

    bnb_model_from_pretrained_args = {}
    if training_args.bits in [4, 8]:
        from transformers import BitsAndBytesConfig
        bnb_model_from_pretrained_args.update(dict(
            device_map={"": training_args.device},
            load_in_4bit=training_args.bits == 4,
            load_in_8bit=training_args.bits == 8,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=training_args.bits == 4,
                load_in_8bit=training_args.bits == 8,
                llm_int8_skip_modules=["mm_projector"],
                llm_int8_threshold=6.0,
                llm_int8_has_fp16_weight=False,
                bnb_4bit_compute_dtype=compute_dtype,
                bnb_4bit_use_double_quant=training_args.double_quant,
                bnb_4bit_quant_type=training_args.quant_type  # {'fp4', 'nf4'}
            )
        ))

    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **asdict(action_head_args))
    if 'paligemma2' in model_args.model_name_or_path:
        cond_dim = config.projection_dim
    else:
        cond_dim = config.hidden_size
    if action_head_args.policy_head_type == 'dit_diffusion_policy':
        config.policy_head_size = action_head_args.policy_head_size
        config.policy_head_config = AutoConfig.for_model(model_type=config.policy_head_type,
                                                       model_size=action_head_args.policy_head_size,
                                                       cond_dim=cond_dim, action_dim=action_head_args.action_dim,
                                                         prediction_horizon=data_args.chunk_size,
                                                       state_dim=action_head_args.state_dim,
                                                         is_tinyvla=model_args.is_tinyvla,
                                                         external_vision_encoder=model_args.external_vision_encoder)
    elif action_head_args.policy_head_type == 'unet_diffusion_policy':
        config.policy_head_config = AutoConfig.for_model(model_type=config.policy_head_type,
                                                       global_cond_dim=cond_dim, action_dim=action_head_args.action_dim,
                                                       state_dim=action_head_args.state_dim,
                                                         is_tinyvla=model_args.is_tinyvla)
    elif action_head_args.policy_head_type == 'gemma_scale_dp_policy':
        config.policy_head_size = action_head_args.policy_head_size
        config.policy_head_config = AutoConfig.for_model(model_type=config.policy_head_type,
                                                       model_size=action_head_args.policy_head_size,
                                                       cond_dim=cond_dim, action_dim=action_head_args.action_dim,
                                                       prediction_horizon=data_args.chunk_size,
                                                       state_dim=action_head_args.state_dim,
                                                       is_tinyvla=model_args.is_tinyvla,
                                                       external_vision_encoder=model_args.external_vision_encoder,
                                                       using_joint_attn=model_args.using_joint_attn)
    else:
        raise NotImplementedError(f"Unsupported policy head type {action_head_args.policy_head_type}")
    # for k,v in asdict(action_head_args).items():
    #     setattr(config, k, v)
    setattr(config.policy_head_config, "input_dim", asdict(action_head_args)['action_dim'])
    setattr(config.policy_head_config, "state_dim", asdict(action_head_args)['state_dim'])

    for k,v in asdict(model_args).items():
        setattr(config, k, v)
    config.llm_loss_weight = training_args.llm_loss_weight

    # todo
    # config.vision_config['image_size_wrist'] = model_args.image_size_wrist

    # config.concat = model_args.concat
    if model_args.is_tinyvla:
        rank0_print(f"{RED} This is TinyVLA, Please Check Both Using_film and Using_xattn equals False:Using_film {model_args.using_film}|Using_xattn {model_args.using_xattn} {RESET}")
        time.sleep(1)
    return model_args, data_args, training_args, action_head_args, config, bnb_model_from_pretrained_args
def train_bc(train_dataset=None, val_dataset=None, model=None, config=None, sampler_params=None, tokenizer=None, processor=None):

    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if config['training_args'].bf16 else torch.float32))
    if config['data_args'].history_images_length > 2:
        rank0_print(f"{RED} Using History and Turn to Video mode.{RESET}")
        video = True
    else:
        video = False
    if 'paligemma' in config['model_args'].model_name_or_path.lower():
        data_collator = PaliGemmaVLADataCollatorForSupervisedDataset(multimodal_processor=processor, computed_type=compute_dtype)

    else:
        data_collator = DexVLADataCollatorForSupervisedDataset(multimodal_processor=processor, computed_type=compute_dtype, tokenizer=tokenizer, video=video)
    # print("data loader test............")
    # from torch.utils.data import DataLoader
    # data_loader = DataLoader(train_dataset, batch_size=config['training_args'].per_device_train_batch_size, collate_fn=data_collator, shuffle=True)
    # for batch in data_loader:
    #     # batch = batch.to('cuda')
    #     # batch = {k:v.to('cuda') for k,v in batch.items()}
    #     for k,v in batch.items():
    #         print(k, v.dtype)
    #     # model(**batch)
    #     # time.sleep(1)
    #     del batch
    #     gc.collect()
    # # exit(0)
    model.config.use_cache = True
    model.config.save_pretrained(config['training_args'].output_dir)
    data_module = dict(train_dataset=train_dataset,
                       data_collator=data_collator,
                       eval_dataset=val_dataset
                       )
    trainer = DexVLATrainer(model=model,
                            tokenizer=tokenizer,
                            args=config['training_args'],
                            sampler_params=sampler_params,
                            **data_module)

    trainer.train(resume_from_checkpoint=config['training_args'].resume_from_checkpoint)

    trainer.save_state()

    model.config.use_cache = True

    if config['training_args'].lora_enable:
        state_dict = ml_utils.get_peft_state_maybe_zero_3(
            model.named_parameters(), config['training_args'].lora_bias
        )
        non_lora_state_dict = ml_utils.get_peft_state_non_lora_maybe_zero_3(
            model.named_parameters(), require_grad_only=False
        )
        if config['training_args'].local_rank == 0 or config['training_args'].local_rank == -1:
            model.config.save_pretrained(config['training_args'].output_dir)
            model.save_pretrained(config['training_args'].output_dir, state_dict=state_dict)
            torch.save(non_lora_state_dict,
                       os.path.join(config['training_args'].output_dir, 'non_lora_trainables.bin'))
    else:
        ml_utils.safe_save_model_for_hf_trainer(trainer=trainer,
                                                  output_dir=config['training_args'].output_dir)



def main(all_config=None, model_config=None):
    set_seed(1)
    # command line parameters
    training_args = all_config['training_args'].__dict__
    # get task parameters
    task_config = TASK_CONFIGS[all_config['data_args'].task_name]
    episode_len = task_config['episode_len']
    camera_names = task_config['camera_names']
    dataset_dir = task_config['dataset_dir']
    name_filter = task_config.get('name_filter', lambda n: True)
    stats_dir = task_config.get('stats_dir', None)
    sample_weights = task_config.get('sample_weights', None)

    all_config['camera_names'] = camera_names
    all_config['episode_len'] = episode_len
    model_config.camera_names = camera_names
    # todo this is pythia's tokenizer not paligemma
    # if 'pythia' in all_config['model_args'].model_name_or_path.lower():
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        all_config['model_args'].model_name_or_path,
    )
    multimodal_processor = AutoProcessor.from_pretrained(all_config['model_args'].model_name_or_path)
    # model = None
    model, data_args = ml_utils.load_model(config=all_config, qwen2_vla_config=model_config, rank0_print=rank0_print, tokenizer=tokenizer)

    if 'paligemma' in all_config['model_args'].model_name_or_path.lower():
        rank0_print(f"{RED} Using PaliGemma as VLA backbone {RESET}")
        image_size = all_config['model_args'].model_name_or_path.split('-')[-1]
        rank0_print(f"{RED} PaliGemma using default and constant Image size{image_size}, omitting SuperParamter:[image_size_stable, image_size_wrist] {RESET}")

        vla_process = PaliGemmaVLAProcess(tokenizer=tokenizer, multimodal_processor=multimodal_processor, data_args=all_config['data_args'])
    else:
        rank0_print(f"{RED} Using Qwen2VL as VLA backbone {RESET}")
        vla_process = DexVLAProcess(tokenizer=tokenizer, multimodal_processor=multimodal_processor, data_args=all_config['data_args'], camera_names=camera_names)

    # train_dataset, val_dataset, stats = load_data(camera_names,
    #                                               all_config['data_args'].chunk_size,
    #                                                 config=all_config,
    #                                                 rank0_print=rank0_print,
    #                                                 policy_class=all_config['action_head_args'].policy_head_type,
    #                                                 llava_pythia_process=vla_process)

    train_dataset, val_dataset, stats, sampler_params = load_data(dataset_dir_l=dataset_dir,
                                                                  name_filter=name_filter,
                                                                  camera_names=camera_names,
                                                                  batch_size_train=all_config['training_args'].per_device_train_batch_size,
                                                                  batch_size_val=all_config['training_args'].per_device_eval_batch_size,
                                                                  chunk_size=all_config['data_args'].chunk_size,
                                                                  skip_mirrored_data=all_config['data_args'].skip_mirrored_data,
                                                                  config=all_config,
                                                                  stats_dir_l=stats_dir,
                                                                  rank0_print=rank0_print,
                                                                  policy_class=all_config['action_head_args'].policy_head_type,
                                                                  sample_weights=sample_weights, train_ratio=0.9999, return_dataset=True, llava_pythia_process=vla_process,
                                                                  action_dim=all_config['action_head_args'].action_dim)



    # exit(0)
    stats_path = os.path.join(all_config['training_args'].output_dir, f'dataset_stats.pkl')
    with open(stats_path, 'wb') as f:
        pickle.dump(stats, f)

    best_ckpt_info = train_bc(train_dataset=train_dataset, model=model, val_dataset=val_dataset, config=all_config, tokenizer=tokenizer, processor=multimodal_processor)
    # save dataset stats
    stats_path = os.path.join(all_config['training_args'].output_dir, f'dataset_stats.pkl')
    with open(stats_path, 'wb') as f:
        pickle.dump(stats, f)


if __name__ == '__main__':
    model_args, data_args, training_args, action_head_args, model_config, bnb_model_from_pretrained_args = parse_param()
    config = {
        'model_args':model_args,
        'data_args':data_args,
        'training_args':training_args,
        'action_head_args':action_head_args,
        'bnb_model_from_pretrained_args':bnb_model_from_pretrained_args
    }

    config_dict = {k:asdict(v) if not isinstance(v, dict) else v for k,v in config.items()}

    ckpt = os.path.join(config['training_args'].output_dir, f"checkpoint-{config['training_args'].save_steps}")
    if os.path.exists(ckpt):
        config['training_args'].resume_from_checkpoint = True
        rank0_print(f"{RED}Resuming Training............{RESET}")
    main(all_config=config, model_config=model_config)
    pass


