# Copyright (c) ModelScope Contributors. All rights reserved.
import json
import numpy as np
import os
from copy import deepcopy

from swift.infer_engine import RequestConfig, TransformersEngine
from swift.ray import RayHelper
from swift.utils import get_logger
from .base import Sampler
from .utils import get_messages_md5, get_reward

logger = get_logger()


@RayHelper.worker(group=['sampler', 'prm', 'orm'])
class VanillaSampler(Sampler):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._prepare_sampler()
        self.caches = self.read_cache()

    @RayHelper.function(group='sampler')
    def _prepare_sampler(self):
        if self.args.sampler_engine == 'transformers':
            _Engine = TransformersEngine
        elif self.args.sampler_engine == 'vllm':
            from swift.infer_engine import VllmEngine
            _Engine = VllmEngine
        elif self.args.sampler_engine == 'lmdeploy':
            from swift.infer_engine import LmdeployEngine
            _Engine = LmdeployEngine
        elif self.args.sampler_engine == 'no':
            _Engine = None
        else:
            raise ValueError(f'Cannot find engine name: {self.args.sampler_engine}')
        self.infer_engine = None
        if _Engine:
            self.infer_engine = _Engine(
                self.args.model, model_type=self.args.model_type, template=self.template, **self.args.engine_kwargs)
            self.infer_engine.strict = False

    @RayHelper.function(group='sampler')
    def read_cache(self):
        cache_files = self.args.cache_files
        caches = {}
        for file in cache_files:
            if not os.path.exists(file):
                logger.warning(f'Cache file does not exist: {file}')
                continue
            with open(file, 'r') as f:
                for line in f.readlines():
                    line = line.strip()
                    if not line:
                        continue

                    content = json.loads(line)
                    uuid = content['id']
                    messages = content['messages']
                    if uuid not in caches:
                        caches[uuid] = {'choices': []}
                    assert messages[-1]['role'] == 'assistant'
                    caches[uuid]['choices'].append(messages[-1]['content'])
        return caches

    @staticmethod
    def convert_data_to_rows(data):
        rows = []
        key = list(data.keys())[0]
        data_len = len(data[key])
        for idx in range(data_len):
            row = {key: data[key][idx] for key in data}
            if row.get('images') and 'bytes' in row['images'][0]:
                row['images'] = [img['path'] for img in row['images']]
            rows.append(row)
        VanillaSampler.check_row_valid(rows)
        return rows

    @staticmethod
    def check_row_valid(rows):
        for row in rows:
            assert not row.get('images') or all([isinstance(img, str) and img for img in row['images']])
            assert not row.get('videos') or all([isinstance(video, str) and video for video in row['videos']])
            assert not row.get('audios') or all([isinstance(audio, str) and audio for audio in row['audios']])

    @RayHelper.function(
        group='sampler',
        dispatch=lambda n, i, data:
        ([{
            'messages': data['messages'][i * len(data['messages']) // n:(i + 1) * len(data['messages']) // n]
        }], {}),
        collect='flatten')
    def generate(self, data):
        resp_all = []
        infer_requests = []
        sent = 0
        rows = self.convert_data_to_rows(data)
        for idx, row in enumerate(rows):
            row = deepcopy(row)
            messages = row['messages']
            uuid = get_messages_md5(row)
            if uuid in self.caches:
                choices = self.caches[uuid]['choices']
                if len(choices) == self.args.num_return_sequences:
                    continue
            if self.args.system:
                if messages[0]['role'] == 'system':
                    messages[0]['content'] = self.args.system
                else:
                    messages.insert(0, {'role': 'system', 'content': self.args.system})
            if messages[-1]['role'] == 'assistant':
                messages = messages[:-1]

            row['messages'] = messages
            infer_request = row
            for i in range(self.args.num_return_sequences):
                infer_requests.append(deepcopy(infer_request))
            sent += 1

        request_config = RequestConfig(
            max_tokens=self.args.max_new_tokens,
            temperature=self.args.temperature,
            top_k=self.args.top_k,
            top_p=self.args.top_p,
        )

        resp_list = []
        if len(infer_requests) > 0:
            resp_list = self.infer_engine.infer(infer_requests, request_config=request_config)

        _cur = 0
        for idx, row in enumerate(rows):
            row = deepcopy(row)
            uuid = get_messages_md5(row)
            if uuid in self.caches:
                choices = self.caches[uuid]['choices']
                if len(choices) == self.args.num_return_sequences:
                    row['choices'] = choices
                    resp_all.append(row)
                    continue

            resps = row
            resps['choices'] = []
            for j in range(self.args.num_return_sequences * _cur, self.args.num_return_sequences * (_cur + 1)):
                if not isinstance(resp_list[j], Exception):
                    resps['choices'].append(resp_list[j].choices[0].message.content)
            if resps['choices']:
                resp_all.append(resps)
            _cur += 1
        return resp_all

    @RayHelper.function(group='orm', dispatch='slice', collect='flatten')
    def get_orm_score(self, infer_requests, ground_truth):
        return get_reward(
            self.orm_model, infer_requests, ground_truths=[ground_truth] * len(infer_requests), threshold=0.0)

    @RayHelper.function(group='prm', dispatch='slice', collect='flatten')
    def get_prm_score(self, infer_requests, ground_truth):
        return get_reward(
            self.prm_model,
            infer_requests,
            ground_truths=[ground_truth] * len(infer_requests),
            threshold=self.args.prm_threshold)

    def do_sample(self, data):
        generated = []
        resp_all = self.generate(data)
        for i, resps in enumerate(resp_all):
            choices = resps['choices']
            messages = resps['messages']
            uuid = get_messages_md5(resps)
            assert messages[-1]['role'] == 'assistant'
            ground_truth = messages[-1]['content']

            infer_requests = []
            for decoded in choices:
                _resps = deepcopy(resps)
                _resps['messages'][-1]['content'] = decoded
                infer_requests.append(_resps)

            _resps = deepcopy(resps)
            _resps['messages'][-1]['content'] = ground_truth
            infer_requests.append(_resps)
            if self.args.orm_model is not None:
                orm_score, _orm_mask = self.get_orm_score(infer_requests, ground_truth)
            else:
                orm_score = np.array([1.0] * len(infer_requests))
                _orm_mask = np.array([True] * len(infer_requests))
            if self.args.prm_model is not None:
                prm_score, _prm_mask = self.get_prm_score(infer_requests, ground_truth)
            else:
                prm_score = np.array([1.0] * len(infer_requests))
                _prm_mask = np.array([True] * len(infer_requests))

            _mask = _orm_mask & _prm_mask
            if not any(_mask):
                continue

            choices.append(ground_truth)
            choices = np.array(choices)

            if self.args.orm_model is None and self.args.prm_model is None:
                positives = choices[:-1]
                for positive in positives:
                    _resps = deepcopy(resps)
                    _resps.pop('choices', None)
                    _resps['id'] = uuid
                    _resps['messages'][-1]['content'] = str(positive)
                    generated.append(json.dumps(_resps, ensure_ascii=False) + '\n')
            else:
                score = np.array(prm_score) + np.array(orm_score * 10)
                sorted_indices = np.argsort(score)[::-1]
                pos_indexes = sorted_indices[0:self.args.n_best_to_keep]
                neg_index = sorted_indices[-1]
                pos_indexes = [int(i) for i in pos_indexes if _mask[i] and i != neg_index]
                logger.info(
                    f'orm:{orm_score}, prm:{prm_score}, positive index: {pos_indexes}, negative index: {neg_index}')
                if self.args.easy_query_threshold is not None and sum([score > 0 for score in orm_score]) - 1 >= int(
                        self.args.num_return_sequences * self.args.easy_query_threshold):
                    continue
                if len(pos_indexes) > 0:
                    positives = choices[pos_indexes]
                    negative = choices[neg_index]
                    for positive in positives:
                        _resps = deepcopy(resps)
                        messages = deepcopy(messages)
                        _resps.pop('choices', None)
                        _resps['messages'][-1]['content'] = str(positive)
                        _resps['rejected_response'] = str(negative)
                        _resps['id'] = uuid
                        generated.append(json.dumps(_resps, ensure_ascii=False) + '\n')
        return generated
