# Generates positive movie reviews by tuning a pretrained model on IMDB dataset
# with a sentiment reward function
import json
import os
import sys
from typing import List

import torch
from datasets import load_dataset
from transformers import pipeline

import trlx
from trlx.data.default_configs import TRLConfig, default_ppo_config


def get_positive_score(scores):
    "Extract value associated with a positive sentiment from pipeline's output"
    return dict(map(lambda x: tuple(x.values()), scores))["POSITIVE"]


def get_negative_score(scores):
    return dict(map(lambda x: tuple(x.values()), scores))["NEGATIVE"]


def main(hparams={}):
    # Merge sweep config with default config if given
    config = TRLConfig.update(default_ppo_config().to_dict(), hparams)

    if torch.cuda.is_available():
        device = int(os.environ.get("LOCAL_RANK", 0))
    else:
        device = -1

    sentiment_fn = pipeline(
        "sentiment-analysis",
        "lvwerra/distilbert-imdb",
        top_k=2,
        truncation=True,
        batch_size=256,
        device=device,
    )

    def dense_reward_fn(samples: List[str], prompts: List[str], outputs: List[str], tokenizer, **kwargs) -> List[float]:
        # Reward positively for initially negative then positive review
        # Reward functions should never receive padded text except for a single EOS at the end
        # Reward function should return token rewards for just the response
        first_halves = [".".join(sample.split(".")[: len(sample.split(".")) // 2]) for sample in samples]
        negative_first_halves = list(map(get_negative_score, sentiment_fn(first_halves)))
        second_halves = [".".join(sample.split(".")[len(sample.split(".")) // 2 :]) for sample in samples]
        positive_second_halves = list(map(get_positive_score, sentiment_fn(second_halves)))
        text_scores = [[f, s] for f, s in zip(negative_first_halves, positive_second_halves)]
        tok_scores = []
        for sample, prompt, response, text_score in zip(samples, prompts, outputs, text_scores):
            toks = tokenizer(response).input_ids
            tok_score = [0] * len(toks)
            tok_score[len(tok_score) // 2] = text_score[0]
            tok_score[-1] = text_score[1]
            tok_scores.append(tok_score)
        return tok_scores

    # Take few words off of movies reviews as prompts
    imdb = load_dataset("imdb", split="train+test")
    prompts = [" ".join(review.split()[:4]) for review in imdb["text"]]

    trlx.train(
        reward_fn=dense_reward_fn,
        prompts=prompts,
        eval_prompts=["I don't know much about Hungarian underground"] * 256,
        config=config,
    )


if __name__ == "__main__":
    hparams = {} if len(sys.argv) == 1 else json.loads(sys.argv[1])
    main(hparams)