from pymongo import MongoClient
from app.translate_loader import translate
from app.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, k=0.99):
        '''
        :param w: коэфицент неточности веса
        '''
        user_prompts = []
        with open('./resources/static/data/food_requests.txt', 'r', encoding="utf-8") as file:
            for user_prompt in file:
                try:
                    if self.count_entities(user_prompt) == 1:
                        user_prompts.append(user_prompt)
                except ValueError:
                    continue

        with open('./resources/static/data/train1_instruct_ss.jsonl', 'w') as f:
            for i, product in tqdm(enumerate(self.db.find())):

                random_user_prompts = np.random.choice(user_prompts, 4)

                for j, user_prompt in enumerate(random_user_prompts):
                    product_name = product['product_name']
                    weight_rand = round(random.random() * 150 + 50, 5)
                    product_name += f' {weight_rand} g'

                    if (i * 4 + j) % 20 == 0:
                        nutrient = np.random.choice(["proteins", "fats", "carbohydrates"])
                        nutrinent_dict = {'proteins': 'prots', 'fats': 'fats', 'carbohydrates': 'carbs'}
                        nutr_weight = weight_rand * product['macronutra'][nutrinent_dict[nutrient]] / 100
                        product_name += f" {nutr_weight} g of {nutrient}"

                    generated_user_prompt = self.replace_entities(
                        user_prompt,
                        [product_name, ]
                    )
                    if generated_user_prompt == user_prompt:
                        continue

                    product_json = {
                        "messages": [
                            {
                              "role": "user",
                              "content": generated_user_prompt,
                            },
                            {
                                "role": "assistant",
                                "content": json.dumps([
                                    {
                                        "name": product['product_name'],
                                        "measure": "g",
                                        "weight": {
                                            "min": weight_rand * w,
                                            "max": weight_rand / w,
                                        },
                                        "prots_100g": {
                                            "min": product['macronutra']['prots'] * k,
                                            "max": product['macronutra']['prots'] / k,
                                        },
                                        "fats_100g": {
                                            "min": product['macronutra']['fats'] * k,
                                            "max": product['macronutra']['fats'] / k,
                                        },
                                        "carbs_100g": {
                                            "min": product['macronutra']['carbs'] * k,
                                            "max": product['macronutra']['carbs'] / k,
                                        },
                                        "pols_100g": {
                                            "min": product['macronutra']['polyols'] * k,
                                            "max": product['macronutra']['polyols'] / k,
                                        },
                                        "acid_100g": {
                                            "min": product['macronutra']['acids'] * k,
                                            "max": product['macronutra']['acids'] / k,
                                        },
                                        "fiber_100g": {
                                            "min": product['macronutra']['fiber'] * k,
                                            "max": product['macronutra']['fiber'] / k,
                                        },
                                        "alco_100g": {
                                            "min": product['macronutra']['alcohol'] * k,
                                            "max": product['macronutra']['alcohol'] / k,
                                        },
                                        "kcal_100g": {
                                            "min": product['custom_extra']['estimated_mean_energy_kcal'] * k,
                                            "max": product['custom_extra']['estimated_mean_energy_kcal'] / k,
                                        }
                                    },
                                ])
                            }
                        ]
                    }

                    f.write(json.dumps(product_json))
                    f.write("\n")


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

        with open('./resources/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('./static/data/train1_instruct_bs.jsonl', 'w') as f:
            for i, product in tqdm(enumerate(self.db.find())):
                if i > 0 and cur_products_count == num_products:
                    generated_user_prompt = self.replace_entities(
                        user_prompts[cur_user_prompt_idx % len(user_prompts)],
                        product_names
                    )

                    if generated_user_prompt != user_prompts[cur_user_prompt_idx % len(user_prompts)]:
                        if len(generated_user_prompt) > 0:
                            product_json['messages'][0]['content'] = generated_user_prompt
                            product_json['messages'][1]['content'] = json.dumps(product_json['messages'][1]['content'])

                            f.write(json.dumps(product_json))
                            f.write('\n')

                    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": []
                            }
                        ]
                    }

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

                if i % 20 == 0:
                    nutrient = np.random.choice(["proteins", "fats", "carbohydrates"])
                    nutrient_dict = {'proteins': 'prots', 'fats': 'fats', 'carbohydrates': 'carbs'}
                    nutr_weight = product['macronutra'][nutrient_dict[nutrient]] * weight_rand / 100

                    product_names.append(
                        f"{product['product_name']} {nutr_weight} g of {nutrient}"
                    )
                else:
                    product_names.append(product['product_name'])

                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'] * k,
                            "max": product['macronutra']['prots'] / k,
                        },
                        "fats_100g": {
                            "min": product['macronutra']['fats'] * k,
                            "max": product['macronutra']['fats'] / k,
                        },
                        "carbs_100g": {
                            "min": product['macronutra']['carbs'] * k,
                            "max": product['macronutra']['carbs'] / k,
                        },
                        "pols_100g": {
                            "min": product['macronutra']['polyols'] * k,
                            "max": product['macronutra']['polyols'] / k,
                        },
                        "acid_100g": {
                            "min": product['macronutra']['acids'] * k,
                            "max": product['macronutra']['acids'] / k,
                        },
                        "fiber_100g": {
                            "min": product['macronutra']['fiber'] * k,
                            "max": product['macronutra']['fiber'] / k,
                        },
                        "alco_100g": {
                            "min": product['macronutra']['alcohol'] * k,
                            "max": product['macronutra']['alcohol'] / k,
                        },
                        "kcal_100g": {
                            "min": min(
                                product['custom_extra']['energy_kcal'],
                                product['custom_extra']['estimated_mean_energy_kcal']
                            ) * k,
                            "max": max(
                                product['custom_extra']['energy_kcal'],
                                product['custom_extra']['estimated_mean_energy_kcal']
                            ) / k,
                        },
                        "prots": {
                            "min": product['macronutra']['prots'] * (weight_min / 100) * k,
                            "max": product['macronutra']['prots'] * (weight_max / 100) / k,
                        },
                        "fats": {
                            "min": product['macronutra']['fats'] * (weight_min / 100) * k,
                            "max": product['macronutra']['fats'] * (weight_max / 100) / k,
                        },
                        "carbs": {
                            "min": product['macronutra']['carbs'] * (weight_min / 100) * k,
                            "max": product['macronutra']['carbs'] * (weight_max / 100) / k,
                        },
                        "pols": {
                            "min": product['macronutra']['polyols'] * (weight_min / 100) * k,
                            "max": product['macronutra']['polyols'] * (weight_max / 100) / k,
                        },
                        "acid": {
                            "min": product['macronutra']['acids'] * (weight_min / 100) * k,
                            "max": product['macronutra']['acids'] * (weight_max / 100) / k,
                        },
                        "fiber": {
                            "min": product['macronutra']['fiber'] * (weight_min / 100) * k,
                            "max": product['macronutra']['fiber'] * (weight_max / 100) / k,
                        },
                        "alco": {
                            "min": product['macronutra']['alcohol'] * (weight_min / 100) * k,
                            "max": product['macronutra']['alcohol'] * (weight_max / 100) / k,
                        },
                        "kcal": {
                            "min": min(
                                product['custom_extra']['energy_kcal'],
                                product['custom_extra']['estimated_mean_energy_kcal']
                            ) * (weight_min / 100) * k,
                            "max": max(
                                product['custom_extra']['energy_kcal'],
                                product['custom_extra']['estimated_mean_energy_kcal']
                            ) * (weight_max / 100) / k,
                        }
                    }
                )

                if i % 20 == 0:
                    product_json['messages'][1]['content'][-1][nutrient_dict[nutrient]]['min'] = nutr_weight
                    product_json['messages'][1]['content'][-1][nutrient_dict[nutrient]]['max'] = nutr_weight

                cur_products_count += 1


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()