from pymongo import MongoClient
from app.core.translate_loader import translate
from app.core.ner_loader import ner
from tqdm import tqdm
import numpy as np
import argparse
import random
import json


class InstructConverter:
    def __init__(self):
        CONNECTION_STRING = 'mongodb+srv://doadmin:7pfg65eH480k2wl3@db-mongodb-fibert-data-lake-1a57fff3.mongo.ondigitalocean.com/admin?tls=true&authSource=admin&replicaSet=db-mongodb-fibert-data-lake'
        client = MongoClient(CONNECTION_STRING)
        self.db = client['nutrition_database']['products']


    def replace_entities(self, input_string, product_names):
        translated_string = translate(input_string)
        if isinstance(translated_string, list):
            translated_string = translated_string[0]

        ner_entities = ner(translated_string)
        if len(product_names) < len(ner_entities):
            return ""

        for i, entity in enumerate(ner_entities):
            translated_string = translated_string.replace(entity, product_names[i])

        return translated_string


    def count_entities(self, input_string):
        translated_string = translate(input_string)
        ner_entities = ner(translated_string)
        return len(ner_entities)


    def convert_ss(self, w=0.95):
        '''
        :param w: коэфицент неточности веса
        '''
        with open('./static/data/food_requests.txt', 'r', encoding="utf-8") as file:
            all_user_prompts = file.readlines()
        user_prompts = []
        for user_prompt in all_user_prompts:
            if self.count_entities(user_prompt) == 1:
                user_prompts.append(user_prompt)

        cur_user_prompt_idx = 0
        with open('train1_products_ss.jsonl', 'w') as f:
            f.write('[')
            f.write("\u000a")
            for i, product in tqdm(enumerate(self.db.find())):
                if i > 0:
                    f.write(json.dumps(product_json, indent=4) + ',')
                    f.write("\u000a")
                    cur_user_prompt_idx += 1

                generated_user_prompt = self.replace_entities(user_prompts[cur_user_prompt_idx % len(user_prompts)], product['product_name'])

                weight_rand = random.random() * 150 + 50
                weight_min = weight_rand * w
                weight_max = weight_rand / w

                product_json = {
                    "messages": [
                        {
                          "role": "user",
                          "content": generated_user_prompt,
                        },
                        {
                            "role": "assistant",
                            "content": "{\n" + str([
                                {
                                    "name": product['product_name'],
                                    "measure": "g",
                                    "weight": {
                                        "min": weight_min,
                                        "max": weight_max,
                                    },
                                    "prots_100g": {
                                        "min": product['macronutra']['prots'] * (weight_min / 100),
                                        "max": product['macronutra']['prots'] * (weight_max / 100),
                                    },
                                    "fats_100g": {
                                        "min": product['macronutra']['fats'] * (weight_min / 100),
                                        "max": product['macronutra']['fats'] * (weight_max / 100),
                                    },
                                    "carbs_100g": {
                                        "min": product['macronutra']['carbs'] * (weight_min / 100),
                                        "max": product['macronutra']['carbs'] * (weight_max / 100),
                                    },
                                    "pols_100g": {
                                        "min": product['macronutra']['polyols'] * (weight_min / 100),
                                        "max": product['macronutra']['polyols'] * (weight_max / 100),
                                    },
                                    "acid_100g": {
                                        "min": product['macronutra']['acids'] * (weight_min / 100),
                                        "max": product['macronutra']['acids'] * (weight_max / 100),
                                    },
                                    "fiber_100g": {
                                        "min": product['macronutra']['fiber'] * (weight_min / 100),
                                        "max": product['macronutra']['fiber'] * (weight_max / 100),
                                    },
                                    "alco_100g": {
                                        "min": product['macronutra']['alcohol'] * (weight_min / 100),
                                        "max": product['macronutra']['alcohol'] * (weight_max / 100),
                                    },
                                    "kcal_100g": {
                                        "min": min(
                                            product['custom_extra']['energy_kcal'],
                                            product['custom_extra']['estimated_mean_energy_kcal']
                                        ) * (weight_min / 100),
                                        "max": max(
                                            product['custom_extra']['energy_kcal'],
                                            product['custom_extra']['estimated_mean_energy_kcal']
                                        ) * (weight_max / 100),
                                    }
                                },
                            ]) + "\n}"
                        }
                    ]
                }

            f.write(json.dumps(product_json, indent=4))
            f.write("\u000a")
            f.write(']')
            f.write("\u000a")


    def convert_bs(self, w=0.95, k=0.99):
        '''
        :param w: коэфицент неточности веса
        :param k: коэфицент неточности нутры
        '''

        with open('./static/data/food_requests.txt', 'r', encoding="utf-8") as file:
            user_prompts = file.readlines()

        cur_user_prompt_idx = 0
        num_products = self.count_entities(user_prompts[cur_user_prompt_idx])
        cur_products_count = 0
        product_names = []

        product_json = {
            "messages": [
                {
                    "role": "user",
                },
                {
                    "role": "assistant",
                    "content": []
                }
            ]
        }

        with open('train1_products_bs.jsonl', 'w') as f:
            f.write('[')
            f.write("\u000a")
            for i, product in tqdm(enumerate(self.db.find())):
                if cur_products_count < num_products:
                    weight_rand = random.random() * 150 + 50
                    weight_min = weight_rand * w
                    weight_max = weight_rand / w

                    product_json['messages'][1]['content'].append(
                        {
                            "name": product['product_name'],
                            "measure": "g",
                            "weight": {
                                "min": weight_min,
                                "max": weight_max,
                            },
                            "prots_100g": {
                                "min": product['macronutra']['prots'] * (weight_min / 100),
                                "max": product['macronutra']['prots'] * (weight_max / 100),
                            },
                            "fats_100g": {
                                "min": product['macronutra']['fats'] * (weight_min / 100),
                                "max": product['macronutra']['fats'] * (weight_max / 100),
                            },
                            "carbs_100g": {
                                "min": product['macronutra']['carbs'] * (weight_min / 100),
                                "max": product['macronutra']['carbs'] * (weight_max / 100),
                            },
                            "pols_100g": {
                                "min": product['macronutra']['polyols'] * (weight_min / 100),
                                "max": product['macronutra']['polyols'] * (weight_max / 100),
                            },
                            "acid_100g": {
                                "min": product['macronutra']['acids'] * (weight_min / 100),
                                "max": product['macronutra']['acids'] * (weight_max / 100),
                            },
                            "fiber_100g": {
                                "min": product['macronutra']['fiber'] * (weight_min / 100),
                                "max": product['macronutra']['fiber'] * (weight_max / 100),
                            },
                            "alco_100g": {
                                "min": product['macronutra']['alcohol'] * (weight_min / 100),
                                "max": product['macronutra']['alcohol'] * (weight_max / 100),
                            },
                            "kcal_100g": {
                                "min": min(
                                    product['custom_extra']['energy_kcal'],
                                    product['custom_extra']['estimated_mean_energy_kcal']
                                ) * (weight_min / 100),
                                "max": max(
                                    product['custom_extra']['energy_kcal'],
                                    product['custom_extra']['estimated_mean_energy_kcal']
                                ) * (weight_max / 100),
                            }
                        }
                    )
                    if i % 20 == 0:
                        nutrient = np.random.choice(["proteins", "fats", "carbohydrates"])
                        nutrinent_dict = {'proteins': 'prots', 'fats': 'fats', 'carbohydrates': 'carbs'}
                        product_names.append(
                            f"{product['product_name']} {product['macronutra'][nutrinent_dict[nutrient]]} g of {nutrient}"
                        )
                    else:
                        product_names.append(product['product_name'])
                    cur_products_count += 1
                else:
                    generated_user_prompt = self.replace_entities(user_prompts[cur_user_prompt_idx % len(user_prompts)], product_names)

                    if len(generated_user_prompt) > 0:
                        product_json['messages'][0]['content'] = generated_user_prompt
                        product_json['messages'][1]['content'] = "{\n" + str(product_json['messages'][1]['content']) + "\n}"

                        if cur_user_prompt_idx > 0:
                            f.write(',')
                            f.write("\u000a")
                        f.write(json.dumps(product_json, indent=4))

                    cur_user_prompt_idx += 1
                    num_products = self.count_entities(user_prompts[cur_user_prompt_idx % len(user_prompts)])
                    cur_products_count = 0
                    product_names = []

                    product_json = {
                        "messages": [
                            {
                                "role": "user",
                            },
                            {
                                "role": "assistant",
                                "content": []
                            }
                        ]
                    }


            f.write("\u000a")
            f.write(']')
            f.write("\u000a")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('convertion_type', type=str)
    args = parser.parse_args()

    converter = InstructConverter()
    if args.convertion_type == 'ss':
        converter.convert_ss()
    elif args.convertion_type == 'bs':
        converter.convert_bs()

Преобразуй это в юпитер тетрадку с вызовом сначала ss потом bs