# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Using FSDPTrainer
"""
import os
import hydra
import ray
import torch
from transformers import PreTrainedTokenizer, AutoTokenizer

from verl import DataProto
from verl.trainer.ppo.ray_trainer import RayPPOTrainer
from verl.utils.fs import copy_to_local
from tests.e2e.envs.digit_completion import CharTokenizer


def make_reward_function(tokenizer, num_examine):

    def arithmetic_sequence_reward_function(data: DataProto, return_dict: bool = False):
        from tests.e2e.envs.digit_completion.task import compute_reward
        reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32)

        for i in range(data.batch.batch_size[0]):
            data_item = data[i]  # DataProtoItem

            prompt_ids = data_item.batch['prompts']

            prompt_length = prompt_ids.shape[-1]

            # extract raw prompt
            valid_prompt_length = data_item.batch['attention_mask'][:prompt_length].sum()
            valid_prompt_ids = prompt_ids[-valid_prompt_length:]

            # extract response
            response_ids = data_item.batch['responses']
            response_length = response_ids.shape[-1]
            response_mask = data.batch['attention_mask'][i][-response_length:]
            valid_response_length = data_item.batch['attention_mask'][prompt_length:].sum()
            valid_response_ids = response_ids[:valid_response_length]

            # decode
            prompt = tokenizer.decode(valid_prompt_ids)
            response = tokenizer.decode(valid_response_ids)
            # remove bos and eos
            prompt = prompt.replace(tokenizer.sep_token, '')
            response = response.replace(tokenizer.eos_token, '')
            if i < num_examine:
                print(prompt, response)

            reward_output = compute_reward(prompt, response)
            dense_reward = reward_output[0].tolist()
            ground_truth_response = reward_output[1]['ground_truth_response']
            if len(dense_reward) > 0:
                last_reward = dense_reward[-1]
            else:
                if len(ground_truth_response) == 0:
                    last_reward = 1
                else:
                    last_reward = 0

            # pad to response_length
            for _ in range(reward_tensor.shape[-1] - len(dense_reward)):
                dense_reward.append(last_reward)

            dense_reward = torch.as_tensor(dense_reward, dtype=torch.float32, device=reward_tensor.device)
            reward_tensor[i] = dense_reward * response_mask

        if return_dict:
            return {"reward_tensor": reward_tensor}
        else:
            return reward_tensor

    return arithmetic_sequence_reward_function


@hydra.main(config_path='../../../../verl/trainer/config', config_name='ppo_trainer', version_base=None)
def main(config):
    ray.init(
        runtime_env={
            'env_vars': {
                'MEGATRON_USE_CUDA_TIMER': '0',
                'MEGATRON_START_PROCESS_TIMER': 'False',
                'TOKENIZERS_PARALLELISM': 'true',
                'NCCL_DEBUG': 'WARN'
            }
        })

    # print initial config
    from pprint import pprint
    from omegaconf import OmegaConf
    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values

    # print the config
    # print initial config
    print('Config after normalizing batch_size')
    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values

    # download the checkpoint from hdfs
    local_path = copy_to_local(config.actor_rollout_ref.model.path)
    local_path = os.path.expanduser(local_path)
    # instantiate tokenizern
    tokenizer = AutoTokenizer.from_pretrained(local_path)
    print(f'Tokenizer vocab_size: {tokenizer.vocab_size}')

    # define worker classes
    from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker
    from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role

    role_worker_mapping = {
        Role.ActorRollout: ray.remote(ActorRolloutRefWorker),
        Role.Critic: ray.remote(CriticWorker),
    }

    global_pool_id = 'global_pool'
    resource_pool_spec = {
        global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
    }
    mapping = {
        Role.ActorRollout: global_pool_id,
        Role.Critic: global_pool_id,
    }

    # use reward model
    if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
        role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)
        mapping[Role.RefPolicy] = global_pool_id

    reward_fn = make_reward_function(tokenizer=tokenizer, num_examine=1)

    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)

    trainer = RayPPOTrainer(config=config,
                            tokenizer=tokenizer,
                            role_worker_mapping=role_worker_mapping,
                            resource_pool_manager=resource_pool_manager,
                            reward_fn=reward_fn,
                            val_reward_fn=reward_fn)
    trainer.init_workers()
    trainer.fit()


if __name__ == '__main__':
    main()