import json
import os
from argparse import ArgumentParser
from typing import Dict, List

from datasets import load_dataset
from transformers import pipeline

import trlx
from trlx.data.default_configs import TRLConfig, default_sft_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 preprocess(instruction: str, input: str, output: str):
    """Build Alpaca prompt and output from instruction and input/output examples"""
    if input:
        prefix = (
            "Below is an instruction that describes a task, paired with an input that provides further context. "
            "Write a response that appropriately completes the request."
        )
        prompt = f"{prefix}\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
        return [prompt, output]
    else:
        prefix = (
            "Below is an instruction that describes a task. Write a response that appropriately completes the request."
        )
        prompt = f"{prefix}\n\n### Instruction:\n{instruction}\n\n### Response:\n"
        return [prompt, output]


def main(hparams={}, model_name="EleutherAI/gpt-j-6B", dataset="tatsu-lab/alpaca"):
    config = default_sft_config()
    config = config.evolve(
        train=dict(
            total_steps=2400,
            batch_size=4,
            seq_length=1024,
        ),
        model=dict(
            model_path=model_name,
        ),
        tokenizer=dict(
            tokenizer_path=model_name,
        ),
        optimizer=dict(kwargs=dict(lr=2e-5)),
        scheduler=dict(kwargs=dict(eta_min=2e-5)),
        method=dict(
            gen_kwargs=dict(
                max_new_tokens=256,
            )
        ),
    )

    # Merge sweep config with default config if given
    config = TRLConfig.update(config.to_dict(), hparams)

    # alpaca = load_dataset("tatsu-lab/alpaca", split="train")
    alpaca = load_dataset(dataset, split="train")
    alpaca = [preprocess(x["instruction"], x["input"], x["output"]) for x in alpaca]

    sentiment_fn = pipeline(
        "sentiment-analysis",
        "lvwerra/distilbert-imdb",
        top_k=2,
        truncation=True,
        batch_size=256,
        device=0 if int(os.environ.get("LOCAL_RANK", 0)) == 0 else -1,
    )

    def metric_fn(samples: List[str], prompts: List[str], outputs: List[str]) -> Dict[str, List[float]]:
        sentiments = list(map(get_positive_score, sentiment_fn(outputs)))
        return {"sentiments": sentiments}

    imdb = load_dataset("imdb", split="test")
    bad_reviews = imdb.filter(lambda sample: sample["label"] == 0).select(range(256))
    zs_rewrite = [preprocess("Rewrite the input into a positive review.", x["text"][:1024], "")[0] for x in bad_reviews]

    trainer = trlx.train(
        samples=alpaca,
        eval_prompts=zs_rewrite,
        metric_fn=metric_fn,
        config=config,
    )

    slug = f"{model_name.split('/')[-1]}-{dataset.split('/')[-1]}"
    trainer.save_pretrained(f"{slug}-sft")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("override_hparams", type=str, default="{}", nargs="?")
    parser.add_argument("--model_name", type=str, default="EleutherAI/gpt-j-6B")
    parser.add_argument("--dataset", type=str, default="tatsu-lab/alpaca")

    args = parser.parse_args()
    hparams = json.loads(args.override_hparams)

    main(hparams, args.model_name, args.dataset)