from flask import Flask, render_template, request, jsonify, send_from_directory
import numpy as np
from PIL import Image
import io
import base64
from scipy.linalg import svd
import os

app = Flask(__name__, static_folder='static', static_url_path='/static')

def calculate_storage(shape, r_value):
    if len(shape) == 3:  # RGB image
        height, width, channels = shape
        original_size = height * width * channels
        compressed_size = channels * (height * r_value + r_value + width * r_value)
    else:  # Grayscale image
        height, width = shape
        original_size = height * width
        compressed_size = height * r_value + r_value + width * r_value
    
    return {
        'original': original_size,
        'compressed': compressed_size,
        'compression_ratio': original_size / compressed_size if compressed_size > 0 else 0
    }

def process_image(image_data, r_value):
    # Open image and convert to grayscale
    img = Image.open(io.BytesIO(image_data))
    
    # Always convert to grayscale
    grayscale_img = img.convert('L')
    img_array = np.array(grayscale_img)
    
    # Perform SVD on the grayscale data
    U, s, Vt = svd(img_array, full_matrices=False)
    
    # Reconstruct image with r singular values
    r = min(r_value, len(s))
    reconstructed = np.dot(U[:, :r] * s[:r], Vt[:r, :])
    
    # Clip values to valid range and convert to uint8
    reconstructed = np.clip(reconstructed, 0, 255).astype(np.uint8)
    
    # Convert back to image
    reconstructed_img = Image.fromarray(reconstructed, mode='L')
    
    # Save to base64 string
    buffered = io.BytesIO()
    reconstructed_img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    
    # Calculate storage requirements
    storage = calculate_storage(img_array.shape, r_value)
    
    return {
        'processed_image': img_str,
        'dimensions': img_array.shape,
        'max_r': min(img_array.shape[0], img_array.shape[1]),
        'storage': storage
    }

def get_predefined_images():
    images_dir = os.path.join('static', 'images')
    if not os.path.exists(images_dir):
        return []
    allowed_extensions = {'.jpg', '.jpeg', '.png', '.gif'}
    images = []
    for filename in os.listdir(images_dir):
        if os.path.splitext(filename)[1].lower() in allowed_extensions:
            images.append({
                'name': os.path.splitext(filename)[0].replace('_', ' ').title(),
                'path': f'/static/images/{filename}'
            })
    return images

@app.route('/')
def index():
    predefined_images = get_predefined_images()
    return render_template('index.html', predefined_images=predefined_images)

@app.route('/process', methods=['POST'])
def process():
    try:
        image_data = None
        if 'image' in request.files:
            image = request.files['image'].read()
            image_data = image
        elif 'image_url' in request.form:
            image_url = request.form['image_url']
            if image_url.startswith('/static/'):
                # Get relative path within static folder
                rel_path = image_url[len('/static/'):]
                image_path = os.path.join(app.static_folder, rel_path)
                
                # Debug info
                app.logger.info(f"Trying to access image at: {image_path}")
                app.logger.info(f"File exists: {os.path.exists(image_path)}")
                
                if os.path.exists(image_path):
                    with open(image_path, 'rb') as f:
                        image_data = f.read()
                else:
                    # Try alternative path resolution
                    alt_path = os.path.join(os.getcwd(), 'static', rel_path)
                    app.logger.info(f"Trying alternative path: {alt_path}")
                    app.logger.info(f"File exists: {os.path.exists(alt_path)}")
                    
                    if os.path.exists(alt_path):
                        with open(alt_path, 'rb') as f:
                            image_data = f.read()
                    else:
                        return jsonify({
                            'error': f'Image file not found. Tried paths: {image_path}, {alt_path}'
                        }), 404

        if not image_data:
            return jsonify({'error': 'No image provided'}), 400

        r_value = int(request.form.get('r_value', 1))
        result = process_image(image_data, r_value)
        return jsonify(result)
    
    except Exception as e:
        import traceback
        app.logger.error(f"Error processing image: {str(e)}")
        app.logger.error(traceback.format_exc())
        return jsonify({
            'error': f'Error processing image: {str(e)}',
            'traceback': traceback.format_exc()
        }), 500
    
if __name__ == '__main__':
    app.run(debug=True)