import viasegura
import pandas as pd
from viasegura import LanesLabeler
from viasegura import ModelLabeler
import numpy as np
import os
import tensorflow as tf
from tqdm.notebook import tqdm
import sys
import csv
import time
import psutil


def calculate_batch_size(sample_image_path, memory_limit_fraction=0.80):
    """
    Calculates the maximum batch size based on available memory.
    :param sample_image_path: Path to a sample image to estimate memory usage
    :param memory_limit_fraction: Fraction of total memory to use (default 75%)
    :return: Maximum batch size
    """
    # Get total available memory
    total_memory = psutil.virtual_memory().total
    memory_limit = total_memory * memory_limit_fraction

    # Estimate memory usage per image
    sample_image = tf.image.decode_image(tf.io.read_file(sample_image_path)).numpy()
    image_memory = sample_image.nbytes  # Bytes per image

    # Calculate maximum batch size
    max_batch_size = int(memory_limit / image_memory)
    print(f"Estimated memory per image: {image_memory / (1024 * 1024):.2f} MB")
    print(f"Available memory for processing: {memory_limit / (1024 * 1024):.2f} MB")
    print(f"Calculated batch size: {max_batch_size}")
    return max_batch_size

def csv_to_filenames_array(csv_file_path):
    """
    Reads a CSV file and returns a list containing all filenames in the order they appear.
    :param csv_file_path: Path to the CSV file
    :return: List of filenames
    """
    all_filenames = []
    with open(csv_file_path, 'r') as file:
        reader = csv.reader(file)
        for row in reader:
            all_filenames.extend(row)
    return all_filenames

def load_image(routes):
    """
    Loads images into memory from the given file paths and calculates memory usage.
    :param routes: List of image file paths
    :return: Numpy array of loaded images and memory usage in MB
    """
    images = np.array([tf.image.decode_image(tf.io.read_file(route)).numpy() for route in routes])
    # Calculate memory usage
    memory_usage = images.nbytes / (1024 * 1024)  # Convert bytes to MB
    print(f"Memory usage of loaded images: {memory_usage:.2f} MB")
    return images

def get_project_path(base_path, cam_choice):
    """
    Builds the project path based on the user's camera choice.
    :param base_path: The base project path
    :param cam_choice: Integer value representing the camera choice (0 for CAM 1, non-0 for CAM 2)
    :return: Full path to the camera directory
    """
    if cam_choice == 0:
        return os.path.join(base_path, "CAM 1")
    else:
        return os.path.join(base_path, "CAM 2")

def process_images(csv_file_path, base_path, cam_choice, start, end):
    """
    Main function to process a subset of images: load filenames from a CSV, generate full paths, and load images.
    :param csv_file_path: Path to the CSV file with image filenames
    :param base_path: Base project path
    :param cam_choice: Integer to select camera directory (0 for CAM 1, non-0 for CAM 2)
    :param start: Starting index for the subset of files
    :param end: Ending index for the subset of files
    :return: Loaded images as a numpy array
    """
    # Get the project path based on the camera choice
    project_path = get_project_path(base_path, cam_choice)

    # Load filenames from the CSV
    archives_frontal = csv_to_filenames_array(csv_file_path)

    # Select the subset of files for the current batch
    archives_frontal_subset = archives_frontal[start:end]

    # Build full file paths for the images
    archives_frontal_paths = [os.path.join(project_path, item) for item in archives_frontal_subset]

    print(f"Loading images {start} to {end} into memory")
    # Load images into memory
    frontal_images = load_image(archives_frontal_paths)
    print(f"Finished loading images {start} to {end}")
    return frontal_images

    
def separate_attributes(input_string):
    # Define the expected categories
    delineation_attributes = {
        'delineation', 'street_lighting', 'carriageway', 'service_road', 
        'road_condition', 'skid_resistance', 'upgrade_cost', 
        'speed_management', 'bicycle_facility', 'quality_of_curve', 
        'vehicle_parking', 'property_access_points'
    }

    area_attributes = {'area_type', 'land_use'}
    lane_attributes = {'number_of_lanes'}

    # Initialize arrays for each category
    delineation_array = []
    area_array = []
    lane_array = []

    # Split the input string into a list of attributes
    attributes = [attr.strip() for attr in input_string.split(',')]

    # Classify each attribute
    for attribute in attributes:
        if attribute in delineation_attributes:
            delineation_array.append(attribute)
        elif attribute in area_attributes:
            area_array.append(attribute)
        elif attribute in lane_attributes:
            lane_array.append(attribute)

    return delineation_array, area_array, lane_array

def main():
    """
    Main entry point of the script.
    Reads the command line arguments and processes images based on the provided input.
    """
    for arg in sys.argv:
        print(arg)
    print("ViaSegura processing started")
    
    # Command-line arguments
    csv_file_path = sys.argv[1]
    cam_choice = int(sys.argv[3])  # Use int to ensure proper comparison in the logic
    base_path = sys.argv[4]
    BatchSize = int(sys.argv[5])
    frontal_filters, lateral_filters, lanes_filter = separate_attributes(sys.argv[2])

    # Load filenames from the CSV to calculate batch indices
    archives_frontal = csv_to_filenames_array(csv_file_path)
    total_images = len(archives_frontal)
    
    sample_image_path = os.path.join(base_path,"CAM 1" ,archives_frontal[0])  # Assuming valid path
    max_batch_size = calculate_batch_size(sample_image_path)
    # Initialize an empty list to store DataFrames
    dataframes = []

    # Process the images in batches
    for start in range(0, total_images, max_batch_size):
        end = min(start + max_batch_size, total_images)

        # Process frontal images if needed
        if len(frontal_filters) > 0 or len(lanes_filter) > 0:
            frontal_images = process_images(csv_file_path, base_path, cam_choice, start, end)

        # Initialize the ModelLabeler for frontal images
        if len(frontal_filters) > 0:
            frontallabeler = ModelLabeler(model_type='frontal', model_filter=frontal_filters, device='/device:GPU:1')
            frontal_results = frontallabeler.get_labels(frontal_images, batch_size=BatchSize)
            frontal_df = pd.DataFrame(frontal_results['numeric_class'])
            if not frontal_df.empty:
                dataframes.append(frontal_df)

        # Process lanes if needed
        if len(lanes_filter) > 0:
            lanes_labeler = LanesLabeler()
            lanes_results = lanes_labeler.get_labels(frontal_images, batch_size=BatchSize)
            lanes_df = pd.DataFrame(lanes_results['numeric_class'])
            if not lanes_df.empty:
                dataframes.append(lanes_df)
        
        # Process lateral images if needed
        if len(lateral_filters) > 0:
            opposite_cam_choice = 1 if cam_choice == 0 else 0
            frontal_images = process_images(csv_file_path, base_path, opposite_cam_choice, start, end)
            laterallabeler = ModelLabeler(model_type='lateral', model_filter=lateral_filters, device='/device:GPU:1')
            lateral_results = laterallabeler.get_labels(frontal_images, batch_size=BatchSize)
            lateral_df = pd.DataFrame(lateral_results['numeric_class'])
            
            
            if not lateral_df.empty:
                dataframes.append(lateral_df)



    # Concatenate all non-empty DataFrames if any exist
    if dataframes:
        resultsNum_df = pd.concat(dataframes, axis=1)
        csvResultPath = os.path.join(os.getenv('PROGRAMDATA'), "Pavesys", "tempViaSeguraResult.csv")
        resultsNum_df.to_csv(csvResultPath, index=False)

    print("Processing Ended")

if __name__ == "__main__":
    main()