'''
Compute low rank approximation of the feed from your webcam.
Also displays framerate and the relative sizes of the first three
singular values (in grayscale mode)

Controls:
    +/- change rank of the approximation
    c - toggle between color/grayscale
    q - quit

In the main function the parameter 'cam_index' should be set to your webcams
index, this is usually 0 or 1 depending on your machine.

It may take a few seconds to initialize the cam.
'''

import cv2
import numpy as np
from scipy.sparse.linalg import svds
import time


def low_rank_approximation(image, rank=10):
    # Convert the image to float32 for compatibility with svds
    float_image = image.astype(np.float32)

    # Adjust the rank to be at least 1 and at most one less than the smallest image dimension
    rank = max(1, min(rank, min(float_image.shape) - 1))

    # Perform partial SVD
    U, S, Vt = svds(float_image, k=rank)

    # svds returns singular values in ascending order, reverse them
    S = S[::-1]
    U = U[:, ::-1]
    Vt = Vt[::-1, :]

    # Reconstruct the image
    approx_image = np.dot(U, np.dot(np.diag(S), Vt))

    # Ensure the pixel values are within the valid range
    approx_image[approx_image < 0] = 0
    approx_image[approx_image > 255] = 255

    # Convert the image back to uint8 for display
    approx_image = approx_image.astype('uint8')

    return approx_image, S[:3]

def process_frame(frame, rank=5):
    # Assuming a color frame, split into RGB channels
    blue_channel, green_channel, red_channel = cv2.split(frame)

    # Process each channel with SVD and low-rank approximation
    processed_channels = [low_rank_approximation(channel, rank)[0] for channel in [blue_channel, green_channel, red_channel]]

    # Merge the channels back
    return cv2.merge(processed_channels)

def process_frame_grayscale(frame, rank=5):
    # Convert the frame to grayscale
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # Apply low-rank approximation to the grayscale frame
    approx_frame,singular_values = low_rank_approximation(gray_frame, rank)

    return approx_frame,singular_values


def main():
    rank=5
    color = False
    
    # Start the webcam
    cam_index=0 #1 when docked 0 otherwise
    cap = cv2.VideoCapture(cam_index) 

    # Variables for FPS calculation
    frame_count = 0
    start_time = time.time()

    try:
        while True:
            # Capture frame-by-frame
            ret, frame = cap.read()
            #if not ret:
            #    break

            # Resize frame for faster processing (optional)
            frame = cv2.resize(frame, (400, 300))

            # Process the frame
            
            if color:
                processed_frame = process_frame(frame, rank=rank)
            else:
                processed_frame, singular_values = process_frame_grayscale(frame, rank=rank)
                ratio1 = singular_values[0] / singular_values[1] if len(singular_values) > 1 else 0
                ratio2 = singular_values[1] / singular_values[2] if len(singular_values) > 2 else 0

            
            
            display_frame = cv2.resize(processed_frame, (1200, 900))
            
            if not color:            
                display_frame = cv2.cvtColor(display_frame, cv2.COLOR_GRAY2BGR)
                

            # Calculate FPS
            frame_count += 1
            end_time = time.time()
            elapsed_time = end_time - start_time
            if elapsed_time > 0:
                fps = frame_count / elapsed_time
            else:
                fps = 0

            # Display FPS on the frame
            cv2.putText(display_frame, f'rank: {rank}   FPS:{ fps:.0f} FPS', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)

            #If gray, display quotients of sing vals on frame
            if not color:
                cv2.putText(display_frame, f's1/s2: {ratio1:.1f}', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)
                cv2.putText(display_frame, f's2/s3: {ratio2:.1f}', (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)

            # Display the resulting frame
            cv2.imshow('Low-Rank Approximation', display_frame)

            key = cv2.waitKey(1) & 0xFF
            if  key == ord('q'):
                break
            elif key== ord('+'):
                rank+=1
            elif key== ord('-'):
                rank-=1
            elif key== ord('c'):
                color= not color
                
    finally:
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()