from flask import Flask, request, jsonify, render_template, url_for
from flask_socketio import SocketIO
import threading
from ultralytics import YOLO
import numpy as np
import cv2
import matplotlib.pyplot as plt
import importlib
from segment_anything import sam_model_registry, SamPredictor
import os
from werkzeug.utils import secure_filename
import logging
import json
import shutil
import sys
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
app = Flask(__name__)
socketio = SocketIO(app)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
class Config:
    BASE_DIR = os.path.abspath(os.path.dirname(__file__))
    UPLOAD_FOLDER = os.path.join(BASE_DIR, 'static', 'uploads')
    SAM_RESULT_FOLDER = os.path.join(BASE_DIR, 'static', 'sam','sam_results')
    YOLO_RESULT_FOLDER = os.path.join(BASE_DIR, 'static', 'yolo','yolo_results')
    YOLO_TRAIN_IMAGE_FOLDER = os.path.join(BASE_DIR, 'static', 'yolo','dataset_yolo','train','images')
    YOLO_TRAIN_LABEL_FOLDER = os.path.join(BASE_DIR, 'static', 'yolo','dataset_yolo','train','labels')
    AREA_DATA_FOLDER = os.path.join(BASE_DIR, 'static', 'yolo','area_data')
    ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
    MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16MB max file size
    SAM_CHECKPOINT = os.path.join(BASE_DIR, 'static', 'sam',"sam_vit_h_4b8939.pth")
    SAM_2 = os.path.join(BASE_DIR, 'static', 'sam',"sam2.1_hiera_large.pt")
    YOLO_PATH = os.path.join(BASE_DIR, 'static', 'yolo', "model_yolo.pt")
    RETRAINED_MODEL_PATH = os.path.join(BASE_DIR, 'static', 'yolo', "model_retrained.pt")
    DATA_PATH = os.path.join(BASE_DIR, 'static', 'yolo','dataset_yolo', "data.yaml")

app.config.from_object(Config)

# Ensure directories exist
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['SAM_RESULT_FOLDER'], exist_ok=True)
os.makedirs(app.config['YOLO_RESULT_FOLDER'], exist_ok=True)
os.makedirs(app.config['YOLO_TRAIN_IMAGE_FOLDER'], exist_ok=True)
os.makedirs(app.config['YOLO_TRAIN_LABEL_FOLDER'], exist_ok=True)
os.makedirs(app.config['AREA_DATA_FOLDER'], exist_ok=True)


# Initialize Yolo model
try:
    model = YOLO(app.config['YOLO_PATH'])
except Exception as e:
    logger.error(f"Failed to initialize YOLO model: {str(e)}")
    raise

try:
    sam2_checkpoint = app.config['SAM_2']
    model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"

    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
    predictor = SAM2ImagePredictor(sam2_model)
except Exception as e:
    logger.error(f"Failed to initialize SAM model: {str(e)}")
    raise

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def scale_coordinates(coords, original_dims, target_dims):
    """
    Scale coordinates from one dimension space to another.
    
    Args:
        coords: List of [x, y] coordinates
        original_dims: Tuple of (width, height) of original space
        target_dims: Tuple of (width, height) of target space
    
    Returns:
        Scaled coordinates
    """
    scale_x = target_dims[0] / original_dims[0]
    scale_y = target_dims[1] / original_dims[1]
    
    return [
        [int(coord[0] * scale_x), int(coord[1] * scale_y)]
        for coord in coords
    ]

def scale_box(box, original_dims, target_dims):
    """
    Scale bounding box coordinates from one dimension space to another.
    
    Args:
        box: List of [x1, y1, x2, y2] coordinates
        original_dims: Tuple of (width, height) of original space
        target_dims: Tuple of (width, height) of target space
    
    Returns:
        Scaled box coordinates
    """
    scale_x = target_dims[0] / original_dims[0]
    scale_y = target_dims[1] / original_dims[1]
    
    return [
        int(box[0] * scale_x),  # x1
        int(box[1] * scale_y),  # y1
        int(box[2] * scale_x),  # x2
        int(box[3] * scale_y)   # y2
    ]

def retrain_model_fn():
    # Parameters for retraining
    data_path = app.config['DATA_PATH']
    epochs = 5
    img_size = 640
    batch_size = 8

    # Start training with YOLO, using event listeners for epoch completion
    for epoch in range(epochs):
        # Train the model for one epoch, here we simulate with a loop
        model.train(
            data=data_path,
            epochs=1,  # Use 1 epoch per call to get individual progress
            imgsz=img_size,
            batch=batch_size,
            device="cpu"  # Adjust based on system capabilities
        )

        # Emit an update to the client after each epoch
        socketio.emit('training_update', {
            'epoch': epoch + 1,
            'status': f"Epoch {epoch + 1} complete"
        })

    # Emit a message once training is complete
    socketio.emit('training_complete', {'status': "Retraining complete"})
    model.save(app.config['YOLO_PATH'])
    logger.info("Model retrained successfully")

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/yolo')
def yolo():
    return render_template('yolo.html')

@app.route('/upload_sam', methods=['POST'])
def upload_sam_file():
    """
    Handles SAM image upload and embeds the image into the predictor instance.

    Returns:
        JSON response with 'message', 'image_url', 'filename', and 'dimensions' keys
            on success, or 'error' key with an appropriate error message on failure.
    """

    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file part'}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': 'Invalid file type. Allowed types: PNG, JPG, JPEG'}), 400
        
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        # Set the image for predictor right after upload
        image = cv2.imread(filepath)
        if image is None:
            return jsonify({'error': 'Failed to load uploaded image'}), 500
            
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        predictor.set_image(image)
        logger.info("Image embedded successfully")
        
        # Get image dimensions
        height, width = image.shape[:2]
        
        image_url = url_for('static', filename=f'uploads/{filename}')
        logger.info(f"File uploaded successfully: {filepath}")
        
        return jsonify({
            'message': 'File uploaded successfully',
            'image_url': image_url,
            'filename': filename,
            'dimensions': {
                'width': width,
                'height': height
            }
        })
        
    except Exception as e:
        logger.error(f"Upload error: {str(e)}")
        return jsonify({'error': 'Server error during upload'}), 500

@app.route('/upload_yolo', methods=['POST'])
def upload_yolo_file():
    """
    Upload a YOLO image file
    
    This endpoint allows a POST request containing a single image file. The file is
    saved to the uploads folder and the image is embedded into the YOLO model.
    
    Returns a JSON response with the following keys:
    - message: a success message
    - image_url: the URL of the uploaded image
    - filename: the name of the uploaded file
    
    If an error occurs, the JSON response will contain an 'error' key with a
    descriptive error message.
    """
    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file part'}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': 'Invalid file type. Allowed types: PNG, JPG, JPEG'}), 400
        
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        
        image_url = url_for('static', filename=f'uploads/{filename}')
        logger.info(f"File uploaded successfully: {filepath}")
        
        return jsonify({
            'message': 'File uploaded successfully',
            'image_url': image_url,
            'filename': filename,
        })
        
    except Exception as e:
        logger.error(f"Upload error: {str(e)}")
        return jsonify({'error': 'Server error during upload'}), 500

@app.route('/generate_mask', methods=['POST'])
def generate_mask():
    """
    Generate a mask for a given image using the YOLO model
    @param data: a JSON object containing the following keys:
        - filename: the name of the image file
        - normalized_void_points: a list of normalized 2D points (x, y) representing the voids
        - normalized_component_boxes: a list of normalized 2D bounding boxes (x, y, w, h) representing the components
    @return: a JSON object containing the following keys:
        - status: a string indicating the status of the request
        - train_image_url: the URL of the saved train image
        - result_path: the URL of the saved result image
    """
    try:
        data = request.json
        normalized_void_points = data.get('void_points', [])
        normalized_component_boxes = data.get('component_boxes', [])
        filename = data.get('filename', '')
        
        if not filename:
            return jsonify({'error': 'No filename provided'}), 400
            
        image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        if not os.path.exists(image_path):
            return jsonify({'error': 'Image file not found'}), 404

        # Read image
        image = cv2.imread(image_path)
        if image is None:
            return jsonify({'error': 'Failed to load image'}), 500
            
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image_height, image_width = image.shape[:2]
        
        # Denormalize coordinates back to pixel values
        void_points = [
            [int(point[0] * image_width), int(point[1] * image_height)]
            for point in normalized_void_points
        ]
        logger.info(f"Void points: {void_points}")
        
        component_boxes = [
            [
                int(box[0] * image_width),
                int(box[1] * image_height),
                int(box[2] * image_width),
                int(box[3] * image_height)
            ]
            for box in normalized_component_boxes
        ]
        logger.info(f"Void points: {void_points}")

        # Create a list to store individual void masks
        void_masks = []
        
        # Process void points one by one
        for point in void_points:
            # Convert point to correct format: [N, 2] array
            point_coord = np.array([[point[0], point[1]]])
            point_label = np.array([1])  # Single label
            
            masks, scores, _ = predictor.predict(
                point_coords=point_coord,
                point_labels=point_label,
                multimask_output=True  # Get multiple masks
            )
            
            if len(masks) > 0:  # Check if any masks were generated
                # Get the mask with highest score
                best_mask_idx = np.argmax(scores)
                void_masks.append(masks[best_mask_idx])
                logger.info(f"Processed void point {point} with score {scores[best_mask_idx]}")

        # Process component boxes
        component_masks = []
        if component_boxes:
            for box in component_boxes:
                # Convert box to correct format: [2, 2] array
                box_np = np.array([[box[0], box[1]], [box[2], box[3]]])
                masks, scores, _ = predictor.predict(
                    box=box_np,
                    multimask_output=True
                )
                if len(masks) > 0:
                    best_mask_idx = np.argmax(scores)
                    component_masks.append(masks[best_mask_idx])
                    logger.info(f"Processed component box {box}")

        # Create visualization with different colors for each void
        combined_image = image.copy()

        # Font settings for labels
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.6
        font_color = (0,0,0)  # White text color
        font_thickness = 1
        background_color = (255, 255, 255)  # White background for text

        # Helper function to get bounding box coordinates
        def get_bounding_box(mask):
            coords = np.column_stack(np.where(mask))
            x_min, y_min = coords.min(axis=0)
            x_max, y_max = coords.max(axis=0)
            return (x_min, y_min, x_max, y_max)
        
        # Helper function to add text with background
        def put_text_with_background(img, text, pos):
            # Calculate text size
            (text_w, text_h), _ = cv2.getTextSize(text, font, font_scale, font_thickness)
            # Define the rectangle coordinates for background
            background_tl = (pos[0], pos[1] - text_h - 2)
            background_br = (pos[0] + text_w, pos[1] + 2)
            # Draw white rectangle as background
            cv2.rectangle(img, background_tl, background_br, background_color, -1)
            # Put the text over the background rectangle
            cv2.putText(img, text, pos, font, font_scale, font_color, font_thickness, cv2.LINE_AA)

        def get_safe_label_position(x_min, y_min, x_max, y_max, text_w, text_h, img_width, img_height):
            # Default to top-right of bounding box
            x_pos = min(y_max, img_width - text_w - 10)  # Keep 10px margin from the right
            y_pos = max(x_min + text_h + 5, text_h + 5)  # Keep 5px margin from the top
            return x_pos, y_pos


        # Apply void masks with different colors
        for mask in void_masks:
            mask = mask.astype(bool)
            combined_image[mask, 0] = np.clip(0.5 * image[mask, 0] + 0.5 * 255, 0, 255)  # Red channel with transparency
            combined_image[mask, 1] = np.clip(0.5 * image[mask, 1], 0, 255)              # Green channel reduced
            combined_image[mask, 2] = np.clip(0.5 * image[mask, 2], 0, 255)
            logger.info("Mask Drawn")  

        # Apply component masks in green
        for mask in component_masks:
            mask = mask.astype(bool)
        # Only apply green where there is no red overlay
            non_red_area = mask & ~np.any([void_mask for void_mask in void_masks], axis=0)
            combined_image[non_red_area, 0] = np.clip(0.5 * image[non_red_area, 0], 0, 255)              # Reduced red channel
            combined_image[non_red_area, 1] = np.clip(0.5 * image[non_red_area, 1] + 0.5 * 255, 0, 255)  # Green channel
            combined_image[non_red_area, 2] = np.clip(0.5 * image[non_red_area, 2], 0, 255)
            logger.info("Mask Drawn") 


        # Add labels on top of masks
        for i,mask in enumerate(void_masks):
            x_min, y_min, x_max, y_max = get_bounding_box(mask)
            (text_w, text_h), _ = cv2.getTextSize("Void", font, font_scale, font_thickness)
            label_position = get_safe_label_position(x_min, y_min, x_max, y_max, text_w, text_h, combined_image.shape[1], combined_image.shape[0])
            put_text_with_background(combined_image, f"Void {i+1}", label_position)    

        for i,mask in enumerate(component_masks):
            i=i+1
            x_min, y_min, x_max, y_max = get_bounding_box(mask)
            (text_w, text_h), _ = cv2.getTextSize("Component", font, font_scale, font_thickness)
            label_position = get_safe_label_position(x_min, y_min, x_max, y_max, text_w, text_h, combined_image.shape[1], combined_image.shape[0])
            put_text_with_background(combined_image, f"Component {i}", label_position)

        # Prepare an empty list to store the output in the required format
        mask_coordinates = []

        for mask in void_masks:
            # Get contours from the mask
            contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            # Image dimensions
            height, width = mask.shape

            # For each contour, extract the normalized coordinates
            for contour in contours:
                contour_points = contour.reshape(-1, 2)  # Flatten to (N, 2) where N is the number of points
                normalized_points = contour_points / [width, height]  # Normalize to (0, 1)

                class_id = 1  # 1 for voids
                row = [class_id] + normalized_points.flatten().tolist()  # Flatten and add the class
                mask_coordinates.append(row)

        for mask in component_masks:
            # Get contours from the mask
            contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            # Filter to keep only the largest contour
            contours = sorted(contours, key=cv2.contourArea, reverse=True)
            largest_contour = [contours[0]] if contours else []
            # Image dimensions
            height, width = mask.shape

            # For each contour, extract the normalized coordinates
            for contour in largest_contour:
                contour_points = contour.reshape(-1, 2)  # Flatten to (N, 2) where N is the number of points
                normalized_points = contour_points / [width, height]  # Normalize to (0, 1)

                class_id = 0  # for components 
                row = [class_id] + normalized_points.flatten().tolist()  # Flatten and add the class
                mask_coordinates.append(row)

        mask_coordinates_filename = f'{filename}.txt'  # Create a unique filename
        mask_coordinates_path = os.path.join(app.config['YOLO_TRAIN_LABEL_FOLDER'], mask_coordinates_filename)


        with open(mask_coordinates_path, "w") as file:
            for row in mask_coordinates:
                # Join elements of the row into a string with spaces in between and write to the file
                file.write(" ".join(map(str, row)) + "\n")

        # Save train image
        train_image_filepath = os.path.join(app.config['YOLO_TRAIN_IMAGE_FOLDER'], filename)
        shutil.copy(image_path, train_image_filepath)
        train_image_url = url_for('static', filename=f'yolo/dataset_yolo/train/images/{filename}')

        # Save result
        result_filename = f'segmented_{filename}'
        result_path = os.path.join(app.config['SAM_RESULT_FOLDER'], result_filename)
        plt.imsave(result_path, combined_image)
        logger.info("Mask generation completed successfully")
        
        return jsonify({
            'status': 'success',
            'train_image_url':train_image_url,
            'result_path': url_for('static', filename=f'sam/sam_results/{result_filename}')
        })

    except Exception as e:
        logger.error(f"Mask generation error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/classify', methods=['POST'])
def classify():
    """
    Classify an image and return the classification result, area data, and the annotated image.

    Request body should contain a JSON object with a single key 'filename' specifying the image file to be classified.

    Returns a JSON object with the following keys:

    - status: 'success' if the classification is successful, 'error' if there is an error.
    - result_path: URL of the annotated image.
    - area_data: a list of dictionaries containing the area and overlap statistics for each component.
    - area_data_path: URL of the JSON file containing the area data.

    If there is an error, returns a JSON object with a single key 'error' containing the error message.
    """

    try:
        data = request.json
        filename = data.get('filename', '')
        if not filename:
            return jsonify({'error': 'No filename provided'}), 400
            
        image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        if not os.path.exists(image_path):
            return jsonify({'error': 'Image file not found'}), 404
        
        # Read image
        image = cv2.imread(image_path)
        if image is None:
            return jsonify({'error': 'Failed to load image'}), 500 

        results = model(image)
        result = results[0]

        component_masks = []
        void_masks = []

        # Extract masks and labels from results
        for mask, label in zip(result.masks.data, result.boxes.cls):
            mask_array = mask.cpu().numpy().astype(bool)  # Convert to a binary mask (boolean array)
            if label == 1:  # Assuming label '1' represents void
                void_masks.append(mask_array)
            elif label == 0:  # Assuming label '0' represents component
                component_masks.append(mask_array)

        # Calculate area and overlap statistics
        area_data = []
        for i, component_mask in enumerate(component_masks):
            component_area = np.sum(component_mask).item()  # Total component area in pixels
            void_area_within_component = 0
            max_void_area_percentage = 0
            
            # Calculate overlap of each void mask with the component mask
            for void_mask in void_masks:
                overlap_area = np.sum(void_mask & component_mask).item()  # Overlapping area
                void_area_within_component += overlap_area
                void_area_percentage = (overlap_area / component_area) * 100 if component_area > 0 else 0
                max_void_area_percentage = max(max_void_area_percentage, void_area_percentage)
            
            # Append data for this component
            area_data.append({
                "Image": filename,
                'Component': f'Component {i+1}',
                'Area': component_area,
                'Void Area (pixels)': void_area_within_component,
                'Void Area %': void_area_within_component / component_area * 100 if component_area > 0 else 0,
                'Max Void Area %': max_void_area_percentage
            })

        area_data_filename = f'area_data_{filename.split("/")[-1]}.json'  # Create a unique filename
        area_data_path = os.path.join(app.config['AREA_DATA_FOLDER'], area_data_filename)

        with open(area_data_path, 'w') as json_file:
            json.dump(area_data, json_file, indent=4)

        annotated_image = result.plot() 

        output_filename = f'output_{filename}'
        output_image_path = os.path.join(app.config['YOLO_RESULT_FOLDER'], output_filename)
        plt.imsave(output_image_path, annotated_image) 
        logger.info("Classification completed successfully")  

        return jsonify({
            'status': 'success',
            'result_path': url_for('static', filename=f'yolo/yolo_results/{output_filename}'),
            'area_data': area_data,
            'area_data_path': url_for('static', filename=f'yolo/area_data/{area_data_filename}')
        })
    except Exception as e:
        logger.error(f"Classification error: {str(e)}")
        return jsonify({'error': str(e)}), 500

retraining_status = {
    'status': 'idle',
    'progress': None,
    'message': None
}

@app.route('/start_retraining', methods=['GET', 'POST'])
def start_retraining():
    """
    Start the model retraining process.

    If the request is a POST, start the model retraining process in a separate thread.
    If the request is a GET, render the retraining page.

    Returns:
        A JSON response with the status of the retraining process, or a rendered HTML page.
    """
    if request.method == 'POST':
        # Reset status
        global retraining_status
        retraining_status['status'] = 'in_progress'
        retraining_status['progress'] = 'Initializing'
        
        # Start retraining in a separate thread
        threading.Thread(target=retrain_model_fn).start()
        return jsonify({'status': 'started'})
    else:
        # GET request - render the retraining page
        return render_template('retrain.html')

# Event handler for client connection
@socketio.on('connect')
def handle_connect():
    print('Client connected')


if __name__ == '__main__':
    app.run(port=5001, debug=True)