

# Strategy: throw everything here, put them into modules if too long or when necessary 







### Introduction 

# find matchups, mean and median for each of the window sizes 
# sensor in the last column 


### Problem

# POLYMER no sensing time for L8, therefore add a day to the time window 
# should be okay because L8 rarely visits the same place two days in a row





import numpy as np
import sys
import glob, os
import pyproj
import datetime as dt
import pandas as pd
import time
import warnings
from importlib.metadata import version
from scipy.spatial import cKDTree


from . import AC_polymer
from . import AC_acolite
from . import AC_l2gen


# make time column and string an option 



def run(sat_folder=None, is_file=None, AC=None, out_file=None, to_log = True, 
        dt_column='GLORIA_time', dt_string = "%Y-%m-%dT%H:%M", output_sd = False,
        ACOLITE_L2R=False, ACOLITE_L1R=False):
    
    
    ### Settings to be included later 
    time_window_minutes = 720 # 60*24 # 720 
    max_dist = 100
    
    
    # make it adaptive to sensor resolution ??
    # window_size = 3
    window_size = [1, 3, 5, 7]
    
    
    ### Print all metadata/settings and save them in a txt file
    
    if to_log: 
        # Start logging in txt file
        orig_stdout = sys.stdout
        
        log_file = out_file.replace(".csv", ".txt")
        # log_file = 'tmart_log_' + time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(time.time())) + '.txt'
    
        class Logger:
            def __init__(self, filename):
                self.console = sys.stdout
                self.file = open(filename, 'w')
                self.file.flush()
            def write(self, message):
                self.console.write(message)
                self.file.write(message)
            def flush(self):
                self.console.flush()
                self.file.flush()
    
        sys.stdout = Logger(log_file)
        
    
    # Metadata
    print('YW_matchups version: ' + str(version('YW_matchups')))
    print('System time: ' + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())))
    print('sat_folder: ' + str(sat_folder))
    print('In-situ file: ' + str(is_file))
    print('AC: ' + str(AC))
    print('out_file: ' + str(out_file))
    print('dt_column: ' + str(dt_column))
    print('dt_string: ' + str(dt_string))
    print('output_sd: ' + str(output_sd))
    print('ACOLITE_L2R: ' + str(ACOLITE_L2R))
    print('ACOLITE_L1R: ' + str(ACOLITE_L1R))
    if to_log: print('log_file: ' + str(log_file))
    
    # Settings 
    print('window_size: ' + str(window_size))
    print('time_window_minutes: ' + str(time_window_minutes))
    print('max_dist: ' + str(max_dist))
    
    
    # Read in-situ data 
    df_is = pd.read_csv(is_file)
    dt_is_string = list(df_is.loc[:,dt_column])
    dt_is_obj = [dt.datetime.strptime(dt_is_string[d],dt_string)for d in list(range(len(dt_is_string)))]
    
    
    
    
    ### find all NC files 
    
    # Keywords for .nc files 
    if AC=='POLYMER' or AC=='L2gen': 
        AC_file = '*.nc'
    elif AC=='ACOLITE':
        AC_file = '*L2W.nc'
        if ACOLITE_L2R or ACOLITE_L1R: AC_file = '*L2R.nc'
    else: 
        sys.exit('Unknown AC algorithm')
    
    # Find them in the directory and subdirectory 
    # https://stackoverflow.com/questions/18394147/how-to-do-a-recursive-sub-folder-search-and-return-files-in-a-list
    sat_files = [y for x in os.walk(sat_folder) for y in glob.glob(os.path.join(x[0], AC_file))]
    sat_files.sort()
    
    
    
    # Output 
    df_output = pd.DataFrame()
    
    
    ### Loop through each of the sat_files
    
    # sat_i = 1
    # sat_file = sat_files[sat_i]
    
    for sat_i, sat_file in enumerate(sat_files): 
        
        print('\n{}/{}'.format(sat_i, len(sat_files)))
        print('Processing: {}'.format(sat_file))
        
        tw_minutes = time_window_minutes
        
        
        
        # Read datetime, for different AC
        if AC=='POLYMER': 
            dt_sat_obj, add_a_day = AC_polymer.get_datetime(sat_file)
        elif AC=='ACOLITE':
            dt_sat_obj, add_a_day = AC_acolite.get_datetime(sat_file)
        elif AC=='L2gen':
            dt_sat_obj, add_a_day = AC_l2gen.get_datetime(sat_file)

        # add a day for POLYMER L8 because we don't have the sensing time 
        # this is okay because L8 doesn't visit the same place two days in a row
        if add_a_day: tw_minutes = tw_minutes + 1440


        # Indices of in-situ data that match the sat datetime  
        is_matches = [ind for ind, s in enumerate(dt_is_obj) if abs((dt_sat_obj - s).total_seconds()/60) <= tw_minutes]
        

        # If there are no matching records, skip to next satellite file.
        if len(is_matches)==0:
            print("No matches for this file within the time limit. Next...\n\n")
            continue
        elif len(is_matches) > 1:
            print('Number of matchups: ' + str(len(is_matches)))

        ### Read NC data, for different AC
        if AC=='POLYMER': 
            sat_data = AC_polymer.read_file(sat_file)
        elif AC=='ACOLITE':
            sat_data = AC_acolite.read_file(sat_file, ACOLITE_L2R, ACOLITE_L1R)
        elif AC=='L2gen':
            sat_data = AC_l2gen.read_file(sat_file)
        
        
        sat_lon = sat_data['lon']
        sat_lat = sat_data['lat']
        
        
        sat_lon_max = sat_data['lon'].max()
        sat_lon_min = sat_data['lon'].min()
        sat_lat_max = sat_data['lat'].max()
        sat_lat_min = sat_data['lat'].min() 
     
        
     
        # Calculate the start time
        start = time.time()
        
        print('Building KDT tree...')
        
        
        ny, nx = sat_lon.shape
        # Flatten the longitude and latitude arrays
        sat_lon_flat = sat_lon.ravel()
        sat_lat_flat = sat_lat.ravel()
        # Stack lon and lat to create an array of points
        points = np.column_stack((sat_lon_flat, sat_lat_flat))
        # Build a KDTree
        tree = cKDTree(points)
        
        
        # Calculate the end time and time taken
        end = time.time()
        length = end - start
        
        # Show the results : this can be altered however you like
        print("Took ", length, " seconds\n")
                

        
        ### Loop through indices and row numbers of in-situ matchups
        
        # ind = 11
        # row = is_matches[ind]
        
        for ind, row in enumerate(is_matches):
            print('Processing in-situ match {}, row {} in spreadsheet'.format(ind, row))
            
            ### Identify distances 
            
            in_situ_lat = df_is.loc[row,"lat"]
            in_situ_lon = df_is.loc[row,"lon"]
            
 
            
            
            # Make sure in situ is in sat range 
            if in_situ_lat > sat_lat_max or in_situ_lat < sat_lat_min:
                print("Not in the range. Next...\n")
                continue
            if in_situ_lon > sat_lon_max or in_situ_lon < sat_lon_min:
                print("Not in the range. Next...\n")
                continue
            
            
                    
            # Query the tree for the closest point
            dist, idx = tree.query([in_situ_lon, in_situ_lat])
            # Convert the flat index back to 2D indices
            line, pixel = np.unravel_index(idx, (ny, nx))
            
            
    

            
            '''
            
            # Project X and Y in meters, centred at the in-situ measurement 
            proj = pyproj.Proj(proj="gnom", ellps="WGS84", lat_0=in_situ_lat,
                               lon_0=in_situ_lon, datum="WGS84", units="m")
            x, y = proj(sat_data['lon'], sat_data['lat'])
            in_situ_x, in_situ_y = proj(in_situ_lon, in_situ_lat)
            
            # Calculate distance and indices of the closest 
            dist = np.sqrt((x - in_situ_x)**2 + (y - in_situ_y)**2)
            
            # Pixels within the max distance, True/False
            within_max_dist = dist < max_dist
            
            # Are there any valid pixels left?
            if np.sum(within_max_dist)==0:
                print("No pixels within acceptable distance (",max_dist,"m ). Next...\n")
                continue
            
            ### Find closest and extract values 
            
            indices_closest = np.argwhere(dist == np.min(dist))
            
            line = indices_closest[0][0]
            pixel = indices_closest[0][1]
            
            '''
            
            
            
            # extract the row as a series 
            row_temp = df_is.iloc[row].copy()
            row_temp['image'] = sat_file
            
            # Loop through window size 
            for ws in window_size: 
            
                # Get pixel coordinates of window around the matching point.
                # (v = vertical/line, h = horizontal/pixel)
                v1 = np.int32(line - ((ws-1)/2))
                v2 = np.int32(line + ((ws+1)/2))
                h1 = np.int32(pixel - ((ws-1)/2))
                h2 = np.int32(pixel + ((ws+1)/2))
                
                # Trim edges of window if they're on the edge of the grid.
                line_start,line_end = max(0, v1), min(sat_data['height'], v2)
                # Window pixel start, pixel end.
                pixel_start,pixel_end = max(0, h1), min(sat_data['width'], h2)
                
                # Extract all values in this window
                rhow_data_window = sat_data['rhow_data'][:,line_start:line_end,pixel_start:pixel_end]
                
                # Column names 
                columns_mean    = ['mean_'+item+'_ws'+str(ws) for item in sat_data['column_names']]
                columns_median  = ['median_'+item+'_ws'+str(ws) for item in sat_data['column_names']]
                
                # SD output 
                if output_sd and ws>1: columns_sd  = ['sd_'+item+'_ws'+str(ws) for item in sat_data['column_names']]
                
                
                # Suppress mean/median of empty slice
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", category=RuntimeWarning)
                
                    # Mean and median values per band 
                    rhow_data_mean      = np.nanmean(rhow_data_window, axis=(1,2))
                    rhow_data_median    = np.nanmedian(rhow_data_window, axis=(1,2))
                    
                    if output_sd and ws>1: rhow_data_sd = np.nanstd(rhow_data_window, axis=(1,2))
                
                # if rhow_data_mean.mask.all(): 
                #     print('All pixels are masked in the window.')
                #     continue
                
                # Add values to data frame
                row_mean    = pd.Series(data=rhow_data_mean,    index=columns_mean)
                row_median  = pd.Series(data=rhow_data_median,  index=columns_median)
                if output_sd and ws>1: row_sd = pd.Series(data=rhow_data_sd,  index=columns_sd)
                
                
                # if output sd 
                if output_sd and ws>1:
                    row_temp = pd.concat([row_temp, row_mean, row_median, row_sd], axis=0)
                else:
                    row_temp = pd.concat([row_temp, row_mean, row_median], axis=0)

            # Add sensor
            row_temp['sensor'] = sat_data['sensor']

            # Add to dataframe             
            df_output = pd.concat([df_output,row_temp.to_frame().T])
            
    # Stop logging 
    if to_log: sys.stdout = orig_stdout
    
    # Output dataframe 
    df_output.to_csv(out_file, index=False)

    return 0
    
    
    
    
    
    
    
    
    
    
    
    
    










