import os
import sys
import argparse
import kp
import cv2
import numpy as np
from mtcnn.mtcnn import MTCNN
import time
import pickle
import json

SCPU_FW_PATH = r"C:\Users\mason\AppData\Local\Kneron_Academy\firmware\KL520\fw_scpu.bin"
NCPU_FW_PATH = r"C:\Users\mason\AppData\Local\Kneron_Academy\firmware\KL520\fw_ncpu.bin"
MODEL_FILE_PATH = 'R34_G369K.nef'
IMAGE_FILE_PATH = 'Chou1.jpg'

def load_image_safe(image_path):
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"[Error] Image file '{image_path}' not found.")
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"[Error] Failed to load image '{image_path}'. Check the file format (must be jpg, png, etc.).")
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img_rgb, img

def landmarks(detector, img_rgb):
    faces = detector.detect_faces(img_rgb)
    if len(faces) == 0:
        raise ValueError("[Error] No faces detected in the image.")
    face = max(faces, key=lambda x: x['confidence'])
    return face['keypoints']

def affine_matrix(lmks, scale=2.5):
    nose = np.array(lmks['nose'], dtype=np.float32)
    left_eye = np.array(lmks['left_eye'], dtype=np.float32)
    right_eye = np.array(lmks['right_eye'], dtype=np.float32)
    eye_width = right_eye - left_eye
    angle = np.arctan2(eye_width[1], eye_width[0])
    center = nose
    alpha = np.cos(angle)
    beta = np.sin(angle)
    w = np.sqrt(np.sum(eye_width**2)) * scale
    m = [[alpha, beta, -alpha * center[0] - beta * center[1] + w * 0.5],
         [-beta, alpha, beta * center[0] - alpha * center[1] + w * 0.5]]
    return np.array(m), (int(w), int(w))

def extract_vector_data(vector):
    """
    Extract data from the inference output object and convert to a standard numpy array.
    Always returns a flattened 1D array regardless of input shape.
    """
    try:
        # For Kneron InferenceFloatNodeOutput specifically
        if 'InferenceFloatNodeOutput' in str(type(vector)):
            # Try to access the data directly
            if hasattr(vector, 'ndarray'):
                data = vector.ndarray
            elif hasattr(vector, 'data'):
                data = vector.data
            elif hasattr(vector, 'content'):
                data = vector.content
            elif hasattr(vector, 'output'):
                data = vector.output
            else:
                # If no direct data attribute, try to access via shape and indexing
                if hasattr(vector, 'shape'):
                    shape = vector.shape
                    if len(shape) == 4 and shape[0] == 1 and shape[2] == 1 and shape[3] == 1:
                        # Common shape for CNN feature vectors: [1, features, 1, 1]
                        # We need to extract each value manually
                        try:
                            data = np.array([vector[0, i, 0, 0] for i in range(shape[1])], dtype=np.float32)
                            return data  # Return early as this is already flat
                        except Exception as e:
                            print(f"Manual extraction failed: {e}")
                            # Continue to other methods
                
                # Last resort - try numpy conversion
                data = np.array(vector, dtype=np.float32)
            
            # Convert to numpy and flatten
            array_data = np.array(data, dtype=np.float32)
            return array_data.flatten()  # Ensure 1D output
        
        # For regular numpy arrays or similar
        if isinstance(vector, np.ndarray):
            return vector.flatten()  # Ensure 1D output
        
        # For list-like objects
        if hasattr(vector, 'tolist'):
            return np.array(vector.tolist(), dtype=np.float32).flatten()
        
        # Generic conversion
        return np.array(vector, dtype=np.float32).flatten()
            
    except Exception as e:
        print(f"Warning: Error converting vector: {e}")
        print(f"Vector type: {type(vector)}")
        if hasattr(vector, 'shape'):
            print(f"Vector shape: {vector.shape}")
            
        # Return a properly shaped array of zeros as fallback
        if hasattr(vector, 'shape') and len(vector.shape) > 0:
            # Find the largest dimension which is likely the feature dimension
            max_dim = max(vector.shape)
            if max_dim > 10:  # Reasonable size for a feature vector
                return np.zeros(max_dim, dtype=np.float32)
        
        # Default fallback size
        return np.zeros(512, dtype=np.float32)
    
def save_vector(vector, file_path, format='numpy', metadata=None):
    """Save a face vector to file using specified format"""
    directory = os.path.dirname(file_path)
    if directory and not os.path.exists(directory):
        os.makedirs(directory)
    
    # Don't print the entire raw vector - it might be a complex object
    print(f"Saving vector of type: {type(vector)}")
    
    # First, try to extract the data into a standard numpy array
    vector_np = extract_vector_data(vector)
    
    # Check for all-zeros vector which indicates extraction failed
    if np.all(vector_np == 0):
        print("WARNING: Extracted vector contains all zeros - extraction likely failed!")
        
        # Add extra debugging before giving up
        print("Attempting emergency extraction methods...")
        
        # Last-ditch effort - try direct attribute access with common names
        for attr_name in ['data', 'array', 'values', 'tensor', 'vector', 'features']:
            if hasattr(vector, attr_name):
                try:
                    data = getattr(vector, attr_name)
                    vector_np = np.array(data, dtype=np.float32)
                    print(f"Emergency extraction via '{attr_name}' attribute succeeded!")
                    break
                except:
                    continue
    
    # Debug information
    print(f"Extracted vector type: {type(vector_np)}")
    print(f"Extracted vector shape: {vector_np.shape}")
    print(f"Sample values: {vector_np[:5]} ... {vector_np[-5:] if len(vector_np) > 5 else []}")
    
    # Check if the shape needs to be adjusted
    if len(vector_np.shape) > 1:
        vector_np = vector_np.squeeze()
        print(f"Squeezed vector shape: {vector_np.shape}")
    
    # Save according to format
    try:
        if format == 'numpy':
            np.save(file_path, vector_np)
        elif format == 'pickle':
            with open(file_path, 'wb') as f:
                pickle.dump(vector_np, f)
        elif format == 'json':
            data = {
                'vector': vector_np.tolist(),
                'metadata': metadata or {}
            }
            with open(file_path, 'w') as f:
                json.dump(data, f)
        else:
            raise ValueError(f"Unsupported format: {format}")
        
        print(f"Vector saved to {file_path}")
        return file_path
    except Exception as e:
        print(f"Error saving vector: {e}")
        
        # Alternative save method if standard methods fail
        if format == 'numpy' or format == 'pickle':
            # Try saving as JSON as a fallback
            try:
                alt_path = file_path + '.json'
                data = {
                    'vector': vector_np.tolist(),
                    'metadata': metadata or {}
                }
                with open(alt_path, 'w') as f:
                    json.dump(data, f)
                print(f"Vector saved using alternative method to {alt_path}")
                return alt_path
            except Exception as e2:
                print(f"Alternative save method also failed: {e2}")
                return None
        return None

def load_vector(file_path, format='numpy'):
    """Load a face vector from file using specified format"""
    if not os.path.isfile(file_path):
        raise FileNotFoundError(f"Vector file '{file_path}' not found.")
        
    if format == 'numpy' or file_path.endswith('.npy'):
        return np.load(file_path)
    elif format == 'pickle' or file_path.endswith('.pkl'):
        with open(file_path, 'rb') as f:
            return pickle.load(f)
    elif format == 'json' or file_path.endswith('.json'):
        with open(file_path, 'r') as f:
            data = json.load(f)
        return np.array(data['vector']), data.get('metadata')
    else:
        raise ValueError(f"Unsupported format: {format}")

def visualize_alignment(original_img, lmks, aligned_img, save_path='alignment_visualization.jpg'):
    """
    視覺化MTCNN檢測到的特徵點和對齊後的結果
    
    Args:
        original_img: 原始BGR圖像
        lmks: MTCNN檢測到的特徵點字典 ('left_eye', 'right_eye', 'nose', etc.)
        aligned_img: 對齊後的人臉圖像
        save_path: 保存視覺化結果的路徑
    """
    import matplotlib.pyplot as plt
    
    # 創建原始圖像的副本用於繪製
    img_vis = original_img.copy()
    
    # 在原始圖像上繪製特徵點
    cv2.circle(img_vis, tuple(map(int, lmks['left_eye'])), 5, (0, 255, 0), -1)
    cv2.circle(img_vis, tuple(map(int, lmks['right_eye'])), 5, (0, 255, 0), -1)
    cv2.circle(img_vis, tuple(map(int, lmks['nose'])), 5, (0, 255, 0), -1)
    cv2.circle(img_vis, tuple(map(int, lmks['mouth_left'])), 5, (0, 255, 0), -1)
    cv2.circle(img_vis, tuple(map(int, lmks['mouth_right'])), 5, (0, 255, 0), -1)
    
    # 在眼睛之間繪製線條，顯示對齊參考線
    cv2.line(img_vis, 
             tuple(map(int, lmks['left_eye'])), 
             tuple(map(int, lmks['right_eye'])), 
             (255, 0, 0), 2)
    
    # 創建一個圖形來顯示兩張圖像
    plt.figure(figsize=(12, 6))
    
    # 顯示帶有特徵點的原始圖像
    plt.subplot(1, 2, 1)
    plt.imshow(cv2.cvtColor(img_vis, cv2.COLOR_BGR2RGB))
    plt.title('原始圖像與特徵點')
    plt.axis('off')
    
    # 顯示對齊後的圖像
    plt.subplot(1, 2, 2)
    plt.imshow(cv2.cvtColor(aligned_img, cv2.COLOR_BGR2RGB))
    plt.title('對齊後的人臉')
    plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(save_path)
    plt.show()
    
    print(f"視覺化結果已保存到 '{save_path}'")


def cosine_similarity(vec1, vec2):
    """
    Calculate cosine similarity between two vectors.
    Handles different input shapes by flattening both vectors.
    """
    # Ensure both vectors are numpy arrays
    vec1 = np.array(vec1, dtype=np.float32)
    vec2 = np.array(vec2, dtype=np.float32)
    
    # Flatten both vectors to ensure 1D
    vec1 = vec1.flatten()
    vec2 = vec2.flatten()
    
    # Check if vectors have compatible sizes
    if vec1.size != vec2.size:
        print(f"Warning: Vector size mismatch: {vec1.size} vs {vec2.size}")
        # Resize shorter vector or truncate longer vector
        if vec1.size < vec2.size:
            vec2 = vec2[:vec1.size]
        else:
            vec1 = vec1[:vec2.size]
    
    # Calculate cosine similarity
    dot_product = np.dot(vec1, vec2)
    norm_a = np.linalg.norm(vec1)
    norm_b = np.linalg.norm(vec2)
    
    # Handle zero division
    if norm_a < 1e-10 or norm_b < 1e-10:
        print("Warning: Vector with near-zero magnitude detected")
        return 0.0
    
    return dot_product / (norm_a * norm_b)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='KL520 ResNet18 model image inference implementation')
    parser.add_argument('-p',
                        '--port_id',
                        help='Using specified port ID for connecting device (Default: port ID of first scanned Kneron device)',
                        default=28,
                        type=int)
    parser.add_argument('-m',
                        '--model',
                        help='Model file path (.nef) (Default: {})'.format(MODEL_FILE_PATH),
                        default=MODEL_FILE_PATH,
                        type=str)
    parser.add_argument('-i',
                        '--img',
                        help='Image file path (Default: {})'.format(IMAGE_FILE_PATH),
                        default=IMAGE_FILE_PATH,
                        type=str)
    parser.add_argument('-o',
                        '--output',
                        help='Output vector file path (Default: output.npy)',
                        default='face_vectors\\output.npy',
                        type=str)
    parser.add_argument('-f',
                        '--format',
                        help='Output format: numpy, pickle, or json (Default: numpy)',
                        default='numpy',
                        choices=['numpy', 'pickle', 'json'],
                        type=str)
    parser.add_argument('-n',
                        '--name',
                        help='Person name for the face vector (for metadata)',
                        default=None,
                        type=str)
    args = parser.parse_args()

    usb_port_id = args.port_id
    MODEL_FILE_PATH = args.model
    IMAGE_FILE_PATH = args.img
    
    """
    connect the device
    """
    try:
        print('[Connect Device]')
        device_group = kp.core.connect_devices(usb_port_ids=[usb_port_id])
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: connect device fail, port ID = \'{}\', error msg: [{}]'.format(usb_port_id,
                                                                                     str(exception)))
        exit(0)

    """
    setting timeout of the usb communication with the device
    """
    print('[Set Device Timeout]')
    kp.core.set_timeout(device_group=device_group, milliseconds=5000)
    print(' - Success')
    
    """
    upload firmware to device
    """
    try:
        print('[Upload Firmware]')
        kp.core.load_firmware_from_file(device_group=device_group,
                                        scpu_fw_path=SCPU_FW_PATH,
                                        ncpu_fw_path=NCPU_FW_PATH)
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: upload firmware failed, error = \'{}\''.format(str(exception)))
        exit(0)
        
    """
    upload model to device
    """
    try:
        print('[Upload Model]')
        model_nef_descriptor = kp.core.load_model_from_file(device_group=device_group,
                                                            file_path=MODEL_FILE_PATH)
        print(' - Success')
    except kp.ApiKPException as exception:
        print('Error: upload model failed, error = \'{}\''.format(str(exception)))
        exit(0)
        
    """
    MTCNN Part
    """
    print('[Process MTCNN]')
    start = time.time()
    # Create MTCNN detector
    detector = MTCNN(device="CPU:0")

    # Load image
    try:
        img_rgb, img_bgr = load_image_safe(IMAGE_FILE_PATH)
        print(f" - Image loaded: {IMAGE_FILE_PATH}")
    except Exception as e:
        print(str(e))
        exit(0)

    # Get landmarks and calculate affine matrix
    try:
        lmks = landmarks(detector, img_rgb)
        mat, size = affine_matrix(lmks)
        print(" - Face landmarks detected")
    except Exception as e:
        print(str(e))
        exit(0)

    # Apply affine transformation
    aligned_img = cv2.warpAffine(img_bgr, mat, size)
    
    end = time.time()
    print(f" - MTCNN processing time: {end - start:.2f} seconds")

    # 添加此行以可視化對齊過程
    visualize_alignment(img_bgr, lmks, aligned_img)
        
    # Convert aligned_img to BGR565 format and resize
    aligned_img_bgr565 = cv2.cvtColor(aligned_img, cv2.COLOR_BGR2BGR565)
    img_bgr565 = cv2.resize(aligned_img_bgr565, (112, 112), interpolation=cv2.INTER_LINEAR)
    print(" - Image aligned and formatted for inference")
    
    """
    prepare generic image inference input descriptor
    """
    generic_inference_input_descriptor = kp.GenericImageInferenceDescriptor(
        model_id=model_nef_descriptor.models[0].id,
        inference_number=0,
        input_node_image_list=[
            kp.GenericInputNodeImage(
                image=img_bgr565,
                image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
                resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
                padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
                normalize_mode=kp.NormalizeMode.KP_NORMALIZE_KNERON
            )
        ]
    )
    
    """
    starting inference work
    """
    print('[Starting Inference Work]')
    try:
        kp.inference.generic_image_inference_send(device_group=device_group,
                                                  generic_inference_input_descriptor=generic_inference_input_descriptor)

        generic_raw_result = kp.inference.generic_image_inference_receive(device_group=device_group)
        print(" - Inference completed successfully")
    except kp.ApiKPException as exception:
        print(' - Error: inference failed, error = {}'.format(exception))
        exit(0)
    
    """
    retrieve inference node output 
    """
    print('[Retrieve Inference Node Output]')
    inf_node_output_list = []
    for node_idx in range(generic_raw_result.header.num_output_node):
        inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(
            node_idx=node_idx,
            generic_raw_result=generic_raw_result,
            channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW
        )
        inf_node_output_list.append(inference_float_node_output)
    print(' - Success')
    
    """
    Process and save the face embedding vector
    """
    # For face recognition models, typically the output is a feature vector
    # Usually, the feature vector is in the first output node
    if len(inf_node_output_list) > 0:
        face_vector = inf_node_output_list[0]
        # print(face_vector)
        print(f"[Face Vector] Original type: {type(face_vector)}")
        print(f"[Face Vector] Shape: {face_vector.shape}")
        
        # Try to examine the vector object
        print("[Face Vector] Available attributes:", [attr for attr in dir(face_vector) if not attr.startswith('__')])
        
        # Create metadata if name is provided
        metadata = None
        if args.name:
            metadata = {
                'name': args.name,
                'image_path': IMAGE_FILE_PATH,
                'timestamp': time.strftime("%Y-%m-%d %H:%M:%S"),
                'dimensions': 512,
            }
        
        output_format = args.format
        output_path = args.output
        # If saving numpy and the file name doesn’t end with .npy, add it
        if output_format == 'numpy' and not output_path.endswith('.npy'):
            output_path += '.npy'
        # Likewise for pickle
        if output_format == 'pickle' and not (output_path.endswith('.pkl') or output_path.endswith('.pickle')):
            output_path += '.pkl'
            
        
        # Save the vector
        output_path = save_vector(
            face_vector, 
            output_path, 
            format=output_format, 
            metadata=metadata
        )
        
        if output_path:
            print(f"[Result] Face embedding vector saved to: {output_path}")
            print(f" - Format: {output_format}")
            if metadata:
                print(f" - Metadata: {metadata}")
        else:
            print("[Error] Failed to save the vector")
    else:
        print("[Error] No output nodes found in the inference result")
        
    # Clean up
    kp.core.disconnect_devices(device_group=device_group)
    print("[Cleanup] Device disconnected")
    
    if output_format == 'numpy':
        loaded = np.load(output_path)
        print("Reload check:", loaded.shape, loaded.dtype)