# -*- coding: utf-8 -*-
"""
Created on Thu Nov 21 12:43:40 2024

@author: chuffard
"""

#!/usr/bin/env python3

import numpy as np
import pandas as pd
import os
import cv2
import glob


def create_circular_mask(h, w, center=None, radius=None):
    """Creates a circular mask for an image."""
    if center is None:  # Use the center of the image
        center = (w * 0.44526, h * 0.47561)
    if radius is None:  # Use the smallest distance between the center and image walls
        radius = w * 0.35539

    Y, X = np.ogrid[:h, :w]
    dist_from_center = np.sqrt((X - center[0])**2 + (Y - center[1])**2)

    mask = dist_from_center <= radius
    return mask


def create_masked_bg_max_img(files, f, num_of_bg_imgs):
    """Create a max pixel image by averaging max pixel values over a range of images."""
    photo_max = None
    count = num_of_bg_imgs
    while count > 0:
        if f >= num_of_bg_imgs:
            photo = cv2.imread(files[f - count])
        else:
            photo = cv2.imread(files[f + count])

        if photo is None:
            print(f"Warning: Unable to read image {files[f]}")
            continue

        photo_hsv = cv2.cvtColor(photo, cv2.COLOR_BGR2HSV)
        photo_value = photo_hsv[:, :, 2]
        photo_value = cv2.normalize(photo_value, None, 0, 255, cv2.NORM_MINMAX)

        if photo_max is None:
            photo_max = photo_value.copy()
        else:
            max_values = np.maximum(photo_max.astype(int), photo_value.astype(int))
            photo_max = max_values

        count -= 1

    mask = create_circular_mask(2050, 2448, center=None, radius=867)
    maskedphotomax = photo_max.copy()
    maskedphotomax[~mask] = 0  # Set pixels outside the circle to black
    return maskedphotomax


def create_masked_image(files, f):
    """Create a masked image by applying a circular mask."""
    masked_img = cv2.imread(files[f])

    if masked_img is None:
        print(f"Warning: Unable to read image {files[f]}")
        return None

    mask = create_circular_mask(2050, 2448, center=None, radius=867)
    masked_img[~mask] = 0  # Set pixels outside the circle to black
    return masked_img


def measure_atn(maskedphotomax, files, f, masked_img):
    """Measure attenuation based on the ratio of red pixel values."""
    # Small epsilon to avoid division by zero
    epsilon = 1e-10
    
    # Ensure the values are non-zero to avoid NaN results in the calculation
    photo_red = masked_img[:, :, 2]  # OpenCV uses BGR, so red is the third channel
    valid_pixels = (maskedphotomax > 0) & (photo_red > 0)  # Filter out invalid pixels

    # Calculate attenuation safely
    attenuation = np.mean(-np.log(photo_red[valid_pixels] / (maskedphotomax[valid_pixels] + epsilon)))

    file_name = os.path.basename(files[f])
    
    atn = (0.981 * attenuation)+0.001
    data = pd.DataFrame([[file_name, attenuation, atn]], columns=['file_name', 'atn_fix', 'atn'])
    
    
    
    return data


# Paths
im_path = 'C:/SES_in'
out_path = 'C:/SES_out'

# Import the files names
fileslist = glob.glob(im_path + "/*.jpg")

# Normalize backslashes in file paths
fileslist = [path.replace('\\', '/') for path in fileslist]

filesdf = pd.DataFrame(fileslist, columns=['filenames'])

# Create output directory if it doesn't exist
if not os.path.isdir(out_path):
    os.mkdir(out_path)

# Filter files with 'Ext' in their names and sort by date
files = filesdf[filesdf['filenames'].str.contains('Ext')]
files['date'] = files['filenames'].str.split('_').str[-1]
files['collect_time'] = files['filenames'].str.split('_').str[-3]
files = files.sort_values(by=['date'])
sorted_filelist = files['filenames'].tolist()

# Identify files that need to be processed based on existing data
if os.path.isfile(os.path.join(out_path, os.path.basename(out_path) + '_atn.csv')):
    with open(os.path.join(out_path, os.path.basename(out_path) + '_atn.csv'), 'r') as data_link:
        last_line = data_link.readlines()[-1]
    
    for x in np.flip(np.arange(0, len(sorted_filelist))):
        if sorted_filelist[x].split('/')[-1] in last_line:
            files_to_be_analyzed = np.arange(x + 1, len(sorted_filelist))
            break
else:
    files_to_be_analyzed = np.arange(0, len(sorted_filelist))

# Initialize an empty DataFrame for new data to be recorded
num_of_bg_imgs = 20
all_atn_data = pd.DataFrame(columns=['file_name', 'atn'])  # Initialize here

# Only process if there are more than twice the number of background images
if len(sorted_filelist) > (num_of_bg_imgs * 2):
    for f in files_to_be_analyzed:
        maskedphotomax = create_masked_bg_max_img(sorted_filelist, f, num_of_bg_imgs)
        masked_img = create_masked_image(sorted_filelist, f)

        if masked_img is not None:
            # Save intermediate images
           # cv2.imwrite(os.path.join(out_path, 'Photomax_' + os.path.basename(sorted_filelist[f])), maskedphotomax)
           # cv2.imwrite(os.path.join(out_path, 'MaskedImg_' + os.path.basename(sorted_filelist[f])), masked_img)
            
            # Measure attenuation and store the results
            atn_data = measure_atn(maskedphotomax, sorted_filelist, f, masked_img)
            
            # Append new data to all_atn_data
            all_atn_data = pd.concat([all_atn_data, atn_data], ignore_index=True)

    # After the loop, save all accumulated data to CSV
    all_atn_data.to_csv(os.path.join(out_path, os.path.basename(out_path) + '_atn.csv'), index=False)
    print(f"Data saved to {os.path.join(out_path, os.path.basename(out_path) + '_atn.csv')}")