#!/usr/bin/env python
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
from sensor_msgs.msg import Image as SensorImage
from lasr_skills import DescribePeople

import cv2_img

r = rospkg.RosPack()

IMG_ROOT = path.join(r.get_path('bsc_receptionist'),
                     'datasets', 'descriptions')


def main():
    rospy.init_node("unit_test_4")

    results = []
    for fn in sorted([a for a in listdir(IMG_ROOT) if a != '.DS_Store'], key=lambda k: int(k.split(',')[0])):
        fnn = fn.split('.')[0]
        id, e_torso_colour, e_hair_colour, e_glasses, e_hat = fnn.split(',')
        expect_glasses = True if e_glasses == 'yes' else False
        expect_hat = True if e_hat == 'yes' else False
        print(f"{id}: {e_torso_colour} {e_hair_colour} {'glasses' if expect_glasses else ''} {'hat' if expect_hat else ''}")

        img = Image.open(path.join(IMG_ROOT, fn))
        img.thumbnail((640, 640))
        msg = cv2_img.pillow_img_to_msg(img)
        pub = rospy.Publisher('/camera/image_raw',
                              SensorImage, queue_size=1, latch=True)
        pub.publish(msg)

        sm = smach.StateMachine(
            outcomes=['failed', 'end'], output_keys=['people'])

        with sm:
            smach.StateMachine.add('DESCRIBE', DescribePeople(), transitions={
                                   'succeeded': 'end', 'failed': 'failed'})

        if sm.execute() == 'end':
            feet = sm.userdata['people'][0]['features']
            has_glasses = False
            has_hat = False
            torso_colour = ''
            hair_colour = ''

            for foot in feet:
                if foot.name == 'torso':
                    torso_colour = foot.colours[0].colour
                elif foot.name == 'hat':
                    has_hat = True
                elif foot.name == 'hair':
                    hair_colour = foot.colours[0].colour
                elif foot.name == 'glasses':
                    has_glasses = True

            results.append({
                "expected": [e_torso_colour, e_hair_colour, expect_glasses, expect_hat],
                "raw": [torso_colour, hair_colour, has_glasses, has_hat],
                "correct_hat": expect_hat == has_hat,
                "correct_glasses": expect_glasses == has_glasses,
                "correct_torso": True if e_torso_colour == torso_colour else (False if torso_colour == '' else "NEEDS_ATTENTION"),
                "correct_hair": True if e_hair_colour == hair_colour else (False if hair_colour == '' else "NEEDS_ATTENTION")
            })
        else:
            print('detection failure!')

        pub.unregister()

    rospy.signal_shutdown("down")
    utils.module_log("Unit 04", "Evaluate Guest Description Generation",
                     "[purple]Test results are ready![/purple]", {
                         "raw": results,
                         "correct_hat_percent": len([0 for result in results if result['correct_hat']]) / len(results),
                         "correct_glasses_percent": len([0 for result in results if result['correct_glasses']]) / len(results),
                     })

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


if __name__ == "__main__":
    import warnings

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        try:
            typer.run(main)
        except Exception as e:
            print(e)
