#!/usr/bin/env python
# before running:
# rosrun bsc_receptionist python -m pip install openai-whisper
# then use:
# rosrun bsc_receptionist python bsc-project/bsc_receptionist/scripts/unit-03-evaluate-speech-system all
import rospy
import smach
import rospkg
import smach_ros
import whisper

from bsc_receptionist import states, data, utils

import typer
from typing import Optional, Union
from typing_extensions import Annotated
from geometry_msgs.msg import Pose, Point, Quaternion
from os import listdir, path
from PIL import Image
import numpy as np
import json
from lasr_rasa.srv import Rasa

import cv2_img

r = rospkg.RosPack()

AUDIO_ROOT = path.join(r.get_path('bsc_receptionist'),
                       'datasets', 'audio clips')
SETS = ['clean', 'busy coffee shop', 'untrained phrases']
MODELS = ['tiny.en', 'small.en', 'medium.en']
# MODELS = ['tiny.en', 'small.en', 'base.en', 'medium.en']


def main(part: str):
    rospy.init_node("unit_test_3")

    done = False

    if part == '1' or part == 'all':
        print('loading models')
        models = [whisper.load_model(name) for name in MODELS]
        print('loaded models!')

        raw_trans = []
        total_t = {}
        success_t = {}
        goals = {}
        rasa_phrases = set([])

        for model in MODELS:
            total_t[model] = {}
            success_t[model] = {}
            goals[model] = {}

        for sett in SETS:
            for model in MODELS:
                total_t[model][sett] = {}
                success_t[model][sett] = {}
                goals[model][sett] = {}

            files = sorted(listdir(path.join(AUDIO_ROOT, sett)))
            for fn in files:
                if fn == '.DS_Store':
                    continue

                expected_phrase = fn.split('(')[0].strip()
                # if expected_phrase != 'my name is charlie':
                #     continue

                # if done:
                #     continue
                # done=True

                print(f'\nExpecting "{expected_phrase}"')
                for idx, model_name in enumerate(MODELS):
                    model = models[idx]
                    result = model.transcribe(
                        path.join(AUDIO_ROOT, sett, fn))["text"]

                    result = result.strip()  # ignore spaces
                    goals[model_name][sett][result] = goals[model_name][sett][result] + \
                        1 if result in goals[model_name][sett] else 1
                    rasa_phrases.add(result)

                    # remove oddities for comparison
                    result = result.lower()
                    if result.endswith('.') or result.endswith('?') or result.endswith('!'):
                        result = result[:-1]

                    result = result.replace("favorite", "favourite")

                    success_t[model_name][sett][expected_phrase] = success_t[model_name][sett][expected_phrase] \
                        if expected_phrase in success_t[model_name][sett] else 0

                    total_t[model_name][sett][expected_phrase] = total_t[model_name][sett][expected_phrase] + \
                        1 if expected_phrase in total_t[model_name][sett] else 1

                    raw_trans.append({
                        "model": model_name,
                        "set": sett,
                        "expected": expected_phrase,
                        "result": result
                    })
                    if result == expected_phrase:
                        print('Success!')
                        success_t[model_name][sett][expected_phrase] = success_t[model_name][sett][expected_phrase] + \
                            1 if expected_phrase in success_t[model_name][sett] else 1
                    else:
                        print('Failure! Got "', result, '"', sep='')
        rasa_phrases = list(rasa_phrases)

    if part == '1' or part == 'all':
        f = open("unit3-cache.json", 'w')
        f.write(json.dumps(
            {"raw_trans": raw_trans,
             "total_t": total_t,
             "success_t": success_t,
             "goals": goals,
             "rasa_phrases": rasa_phrases}
        ))
        f.close()

    if part == '2':
        f = open("unit3-cache.json", 'r')
        a = json.loads(f.read())
        total_t = a["total_t"]
        success_t = a["success_t"]
        goals = a["goals"]
        rasa_phrases = a["rasa_phrases"]
        f.close()

    rasa_results = {}
    if part == '2' or part == 'all':
        rasa_parse = rospy.ServiceProxy(
            f'/lasr_rasa/parse', Rasa)
        for phrase in rasa_phrases:
            print(f"Phrase: {phrase}")
            result = json.loads(rasa_parse(phrase).json_response)
            print(result)
            rasa_results[phrase] = {
                "intent": result["intent"]["name"],
                "conf": result["intent"]["confidence"],
                "entity_value": list(result["entities"].values())[0][0]["value"] if len(list(result["entities"].values())) == 1 else 'INVALID ENTITY'
            }

    results = {
        "raw_trans": raw_trans,
        "transcribe_success_rate": {
            model: {
                sett: {
                    phrase: success_t[model][sett][phrase] / total_count
                    for phrase, total_count in phrases.items()
                } for sett, phrases in setts.items()
            } for model, setts in total_t.items()
        },
        "goals": goals,
        "rasa_results": rasa_results
    }
    utils.module_log("Unit 03", "Evaluate Speech & Intent Recognition",
                     "[purple]Test results are ready![/purple]", results)

    fr = open('results2.json', 'w')
    fr.write(json.dumps(results))
    fr.close()


if __name__ == "__main__":
    import warnings

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        typer.run(main)
