from pathlib import Path
import pandas as pd
import os
os.chdir(Path(__file__).parent.resolve())
from utils import datenum_to_datetime, RegisteredImage, get_site_tide_data
import openpyxl
import scipy.io as sio
from PIL import Image, ImageDraw
from loguru import logger
import argparse


'''
SCRIPT LOGIC

For each site in Images Parent Directory:
    Iterate through years in reverse:
        Iterate through images names (dates) in reverse:
            Check if image name in 'Registered'? 
            If not, tag and save
            If yes, skip to the next site 
            
'''
#-----------------------------------------------------------------------------#
# Get command-line arguments 

parser = argparse.ArgumentParser()
parser.add_argument("--tag_reregistered", action='store_true', required=False)
parser.add_argument("--tag_registered_deleted", action='store_true', required=False)
parser.add_argument("--overwrite", action='store_true', required=False)
parser.add_argument("--site", type=str, required=False)
args = parser.parse_args()

#-----------------------------------------------------------------------------#
# Retrieve coastsnap/images directory set in coastsnap_sites.csv

code_images_dir = str(Path(os.getcwd()).parent)
sites_csv_path = os.path.join(code_images_dir, "coastsnap_sites.csv")
coastsnap_sites_csv = pd.read_csv(sites_csv_path)
images_parent_dir = coastsnap_sites_csv.parent_directory[0] 
images_dir = os.path.join(images_parent_dir, "Images")

site_list = coastsnap_sites_csv['site_name']  # Loop through sites in coastsnap_sites.csv
#site_list = os.listdir(images_dir)  # Loop through all sites in /CoastSnap/Images

#-----------------------------------------------------------------------------#   
# Tag image function

def tag_image(image2tag):
    
    image = Image.open(image2tag.pathname)
    draw = ImageDraw.Draw(image)
    
    # White Text Box 
    if image2tag.large_filename:
        rect_height = image2tag.height/13
    else:
        rect_height = image2tag.height/20
         
    # Create white text box
    draw.rectangle((0, 0, image2tag.width, rect_height), fill='white')

    # Tag image with text
    draw.text((20, rect_height/4),image2tag.tag, font = image2tag.font, fill=(0, 0, 0)) 
    
    #registered_path_wrong = image2tag.pathname[:-4] + '_registered.jpg' 
    #registered_path = registered_path_wrong.replace('Photoshop', 'Registered')
    image = image.convert('RGB')

    #image.save(registered_path)
    image.save(image2tag.registered_pathname)

#-----------------------------------------------------------------------------# 
# Loop Photoshop directory and print reregistered images

def loop_reregistered(site, tag, reregister_all):
    i = 1
    num_2_tag = 0
    tides_df = False
    print("\nSite: " + site)
    photoshop_path = images_dir +'/'+ site + '/Photoshop'
    try: # Check if site contains 'Processed' directory
        years_list = os.listdir(photoshop_path)
        years_list.reverse()
    except: # Site doesn't contain 'Processed'. Go to next site.
        return 
    # Loop through YEARS
    for year in years_list: 
        if not year.isdigit() and len(year) != 4: # Check folder name is in the format "YYYY"
            continue
        year_path = photoshop_path + '/' + year
        image_list = os.listdir(year_path)
        image_list.reverse()
        # Loop through IMAGES in photoshop
        for image_filename in image_list: 

            image_filename_no_suffix = image_filename.replace("_tagged","");
            reregistered_image_path = year_path + '/' + image_filename_no_suffix   
            if os.path.isfile(year_path + '/' + image_filename_no_suffix):
                
                if not tag: 
                    print(image_filename)
                    num_2_tag += 1
                else:
                    registered_year_path = year_path.replace('Photoshop', 'Registered')
                    registered_image_path = registered_year_path + '/' + image_filename_no_suffix[:-4] + '_registered.jpg' 
                    
                    if args.tag_reregistered and ("_tagged" not in image_filename) and not reregister_all:
                        action = input(image_filename +": ")
                        if action == 'y':
                            print("Tagging image\n")
                            pass
                        else:
                            print("Skip\n")
                            continue
                        
                    # Case: Image not tagged, 'overwrite' = True or 'tag reregistered' = True
                    if not os.path.exists(registered_year_path): # Chech that registered/year directory exists
                        os.makedirs(registered_year_path)     # if not, create it
                    
                    pathname = os.path.join(year_path, image_filename)
                    image2tag = RegisteredImage(pathname, image_filename, registered_image_path) # Create image object
                    if i == 1: # Retrieve Tide Data once for each site
                        tides_df = get_site_tide_data(images_parent_dir, site)
                    if isinstance(tides_df, pd.DataFrame): image2tag.get_tide(tides_df)
                    
                    image2tag.get_font()
                    image2tag.get_dimensions()
                    image2tag.create_tag()
                    tag_image(image2tag)
                    i += 1
                    
                    # Add '_tagged' suffix to image in Photoshop
                    new_pathname = pathname[:-4] + '_tagged.jpg'
                    os.rename(pathname, new_pathname) 
                    
    return num_2_tag
    
#-----------------------------------------------------------------------------#
# Main function (called below)

def main(site):
    
    i=0
    site_complete = False # A flag to stop tagging images when found
    tide_data = False
    font = None
    image2tag = None 
    photoshop_path = images_dir +'/'+ site + '/Photoshop'
    try: # Check if site contains 'Processed' directory
        years_list = os.listdir(photoshop_path)
        years_list.reverse()
    except: # Site doesn't contain 'Processed'. Go to next site.
        return 
    
    # Loop through YEARS
    for year in years_list: 
        if not year.isdigit() and len(year) != 4: # Check folder name is in the format "YYYY"
            continue
        if site_complete: break 
        year_path = photoshop_path + '/' + year
        image_list = os.listdir(year_path)
        image_list.reverse()
        
        # Loop through IMAGES in photoshop
        for image_filename in image_list: 
            registered_year_path = year_path.replace('Photoshop', 'Registered')
            image_filename_no_suffix = image_filename.replace("_tagged","");
            registered_image_path = registered_year_path + '/' + image_filename_no_suffix[:-4] + '_registered.jpg' 
            
            if site not in image_filename: # Check the filename has the site in it
                continue
            
            # Case: Image tagged and in 'Registered'
            if os.path.isfile(registered_image_path): 
                if args.tag_registered_deleted: # If 'tag_registered_deleted' = True, continue to next image
                    continue 
                elif args.overwrite:  # If 'overwrite' = True, pass and tag image
                    pass
                else:                           # Stop tagging and go to the next site
                    site_complete = True; 
                    if i == 0:
                        print(site + ": " + str(i) + " images tagged")
                    elif image2tag and image2tag.tide:
                        print(site + ": " + str(i) + " images tagged with tide")
                    else:
                        print(site + ": " + str(i) + " images tagged no tide")                  
                    return
            
                
            # Case: Image not tagged, 'overwrite' = True or 'tag reregistered' = True
            if not os.path.exists(registered_year_path): # Chech that registered/year directory exists
                os.makedirs(registered_year_path)     # if not, create it
            i += 1
            
            pathname = os.path.join(year_path, image_filename)
            image2tag = RegisteredImage(pathname, image_filename, registered_image_path) # Create image object
            if i == 1: # Retrieve Tide Data once for each site
                tides_df = get_site_tide_data(images_parent_dir, site)
            image2tag.get_font()
            image2tag.get_dimensions()
            if isinstance(tides_df, pd.DataFrame): image2tag.get_tide(tides_df)
            image2tag.create_tag()
            
            tag_image(image2tag)
            
            # Add '_tagged' suffix if not already there
            old_tagged_image_path = year_path + '/' + image_filename_no_suffix
            new_tagged_image_path = year_path + '/' + image_filename_no_suffix[:-4] + '_tagged.jpg'
            if os.path.isfile(old_tagged_image_path):
                os.rename(old_tagged_image_path, new_tagged_image_path) 
            
    # Case: All images tagged
    if not site_complete:
        if i == 0:
            print(site + ": " + str(i) + " images tagged")
        elif image2tag.tide:
            print(site + ": " + str(i) + " images tagged with tide")
        else:
            print(site + ": " + str(i) + " images tagged no tide")  

#-----------------------------------------------------------------------------#
# Case: Tag reregistered images

reregister_all = True
tag = False
count = 0
if args.tag_reregistered:
    print("\n   The following images have been reregistered:")
    if args.site:
        count = loop_reregistered(args.site, tag, reregister_all)
    else:
        for site in site_list: # Code to loop over entire images directory: for site in os.listdir(images_dir): 
            x = loop_reregistered(site, tag, reregister_all)
            count = count + x
            
    if count > 0:
        reregister = input("\n\n   Do you want to tag 'all' " + str(count) + " images or just 'some' of them: ")
        
        print("\n")
        tag = True
        if reregister == 'all':
            print("Tagging ALL reregistered images...")
        elif reregister == 'some':
            print("   Tagging SOME reregistered images. For each image:\n\n    - to TAG, type the letter 'y' and hit enter \n    - to SKIP, just hit enter\n")
            reregister_all = False
        else:
            print("   User must type either 'all' or 'some' and then hit enter.")
            exit()
        if args.site:           # Case 1: User specifies a site 
            loop_reregistered(args.site, tag, reregister_all)
        else:                   # Case 2: No site specified. Tag all images
            for site in site_list: # Code to loop over entire images directory: for site in os.listdir(images_dir):
                loop_reregistered(site, tag, reregister_all)

#-----------------------------------------------------------------------------#
# Case: NOT reregistering images

else:
    # Calling the main function with the two cases
    if args.site:           # Case 1: User specifies a site 
        main(args.site)
    else:                   # Case 2: No site specified. Tag all images
        for site in site_list: # Code to loop over entire images directory: for site in os.listdir(images_dir):
            main(site)