"""
Get SUVI data from NODD at a specific cadence
=============================================
Purpose:
   Download SUVI L1b data at a specific cadence from the NOAA Open Data Dissemination (`NODD <https://www.noaa.gov/information-technology/open-data-dissemination>`__) AWS S3 buckets using the boto3 library. 
"""

# %%
# The NODD AWS S3 buckets contain almost all SUVI L1b data for GOES-16 through GOES-19 (some early mission data is not available).
# The data is publicly available at s3://noaa-goes19/ (replace the 19 with 16, 17, or 18 for the earlier missions). 
# 
# The script below is intended to download data in a specific wavelength channel at a requested cadence.
# For example, using this script, you could download a SUVI 171 Å file from GOES-19 at a cadence of one hour,
# instead of having to download all data on the S3 bucket and discarding what you don't need.
#
# How this script operates: the code will create reference timestamps in the requested period and at the
# requested cadence, and then look for the file that is closest to each reference timestamp. If a file is
# found within the defined tolerance, it is accepted, otherwise rejected.

# %%
# Let's first import the necessary Python libraries. The "UNSIGNED" and "Config" from botocore are imported so we
# don't need a specific username/token for the NODD S3 buckets.
__author__ = "cbethge"
import datetime
import numpy as np
import os
import boto3
from botocore import UNSIGNED
from botocore.config import Config

# %%
# Now we define the start and end of the period that the code is supposed to consider.
# The format for those dates is YYYY-MM-DDThh:mm:ss (ISOT without the milliseconds). 
start_date = '2026-05-15T12:20:00'
end_date   = '2026-05-15T23:30:00'

# %%
# .. note:: 
#
#    The code has to parse *all* files in the requested period! A period of about one month can be handled, 
#    but it will be fairly slow. Anything beyond about one month should be broken up into one-month chunks, 
#    e.g. with an enclosing for-loop.

# %%
# Here are some pretty self-explanatory settings: which satellite to get the data from (16, 17, 18, or 19),
# which wavelength (94, 131, 171, 195, 284, or 304), which file ending (.fits for FITS files or .nc for netCDF),
# where to download the data to, the cadence in minutes, and the tolerance in minutes. 
satellite = 19            # which GOES satellite (16, 17, 18, 19)
wavelength = 171          # which SUVI wavelength channel (94, 131, 171, 195, 284, 304)
file_ending = '.fits'     # which file type (.fits or .nc)
outpath = './'            # where to download the data
cadence = 60.             # requested cadence in minutes
tolerance = 4.            # tolerance in minutes

# %%
# How the tolerance works: in this example script, we requested 171 Å data from GOES-19 at a cadence of one hour,
# starting at '2026-05-15T12:20:00'. The reference timestamps created in this case would be:
#
# | 2026-05-15T12:20:00
# | 2026-05-15T13:20:00
# | 2026-05-15T14:20:00
# | 2026-05-15T15:20:00
# |
# | and so on.
# |
# | Let's say the closest file to the second reference timestamp is
# | "OR_SUVI-L1b-Fe171_G19_s20261351318239_e20261351318249_c20261351318474.fits".
# | The start time for this file is "s20261351318239", or "2026-05-15T13:18:23.9" in ISOT format. In
# | other words, the difference to the respective reference timestamp is: minus 1 minute, 36 seconds,
# | and 1 millisecond. Since the tolerance was set to 4 minutes, this file would be accepted. If the
# | tolerance had been set to 1 minute, the file would be rejected.  

# %%
# Now on to some more settings, they are all Boolean (True or False).
# The first one is:
enforce_requested_period = False
# %%
# | This is best explained with another example.
# | The file closest to the first reference timestamp mentioned above is
# | "OR_SUVI-L1b-Fe171_G19_s20261351218237_e20261351218247_c20261351218479.fits".
# | The timestamp for this file is "s20261351218237", or "2026-05-15T12:18:23.7" in ISOT format.
# | Note that this time is *before* the requested period. If enforce_requested_period is set to False, the
# | the file would be accepted anyway. If it is set to True, all files outside of the requested period are
# | rejected.

# %%
# The second True/False setting is:
short_exp = False
# %% 
# By default, the code will download SUVI L1b long exposures. If short_exp is set to True,
# short exposures will be downloaded instead. Note that for the 94 and 131 Å channels, the
# code cannot distinguish between SUVI short exposures and SUVI short flare exposures based
# on their filenames, so it will download whichever of the two is closest in time to the
# reference timestamp in this case!

# %%
# Moving on with the True/False settings:
no_sub_dirs = False
# %%
# By default, the code will create subdirectories for satellite/year/month/day/wavelength/.
# If no_sub_dirs is set to True, no subdirectories will be created and the files are downloaded
# directly to the "outpath" folder defined above.

# %%
# And finally:
dryrun = True
# %%
# If dryrun is set to True, the filenames of the accepted files will only be printed out, but the
# files will not be downloaded. This is the default setting for this website example script. 

# %%
# Now that we are done with the settings, we do some really basic sanity checks.
# First, we define a function to make sure the submitted dates are in the right format:
def is_valid_date_string(date_string):
    # Defining the format from above for datetime
    date_format = "%Y-%m-%dT%H:%M:%S"

    try:
        # Attempt to parse the string
        datetime.datetime.strptime(date_string, date_format)
        return True
    except ValueError:
        # If a ValueError is raised, the format or the date itself is invalid
        return False

# %%
# Check that the settings from above are valid inputs:
valid_satellites = [16, 17, 18, 19]
assert satellite in valid_satellites,  f'satellite must be in {valid_satellites}'

valid_wavelengths = [94, 131, 171, 195, 284, 304]
assert wavelength in valid_wavelengths,  f'wavelength must be in {valid_wavelengths}'

assert is_valid_date_string(start_date), 'Start date must have the form YYYY-MM-DDThh:mm:ss'

syear   = int(start_date[0:4])
smonth  = int(start_date[5:7])
sday    = int(start_date[8:10])
shour   = int(start_date[11:13])
sminute = int(start_date[14:16])
ssecond = int(start_date[17:19])

assert syear >= 2019,      'Start year must be >= 2017.'
assert 1 <= smonth <= 12,  'Start month must be between 1 and 12.'
assert 1 <= sday <= 31,    'Start day must be between 1 and 31.'
assert 0 <= shour <= 23,   'Start hour must be between 0 and 23.'
assert 0 <= sminute <= 59, 'Start minute must be between 0 and 59.'
assert 0 <= ssecond <= 59, 'Start second must be between 0 and 59.'

start_date = datetime.datetime.fromisoformat(start_date).replace(tzinfo=datetime.timezone.utc)

assert is_valid_date_string(end_date), 'End date must have the form YYYY-MM-DDThh:mm:ss'

eyear   = int(end_date[0:4])
emonth  = int(end_date[5:7])
eday    = int(end_date[8:10])
ehour   = int(end_date[11:13])
eminute = int(end_date[14:16])
esecond = int(end_date[17:19])

assert eyear >= 2017,      'End year must be >= 2017.'
assert 1 <= emonth <= 12,  'End month must be between 1 and 12.'
assert 1 <= eday <= 31,    'End day must be between 1 and 31.'
assert 0 <= ehour <= 23,   'End hour must be between 0 and 23.'
assert 0 <= eminute <= 59, 'End minute must be between 0 and 59.'
assert 0 <= esecond <= 59, 'End second must be between 0 and 59.'

end_date = datetime.datetime.fromisoformat(end_date).replace(tzinfo=datetime.timezone.utc)
assert (end_date-start_date).total_seconds() > 0, 'End date must be after start date.'

# %%
# The main algorithm: 

# Start processing...
satellite_name = 'G'+str(satellite)
satellite_name_long = 'goes'+str(satellite)

# Initialize connection with S3 bucket
suvi_session = boto3.session.Session()
s3 = boto3.resource('s3', config=Config(signature_version=UNSIGNED))
suvi_bucket = s3.Bucket('noaa-'+satellite_name_long)

# These are the SUVI wavelength subfolders on the NODD S3 buckets
wl_folders = { 94: 'SUVI-L1b-Fe093', 131: 'SUVI-L1b-Fe131', 171: 'SUVI-L1b-Fe171',
              195: 'SUVI-L1b-Fe195', 284: 'SUVI-L1b-Fe284', 304: 'SUVI-L1b-He303'}

# We have to parse through all days being considered to make a list of all
# available files on the S3
print('Making list of files on S3 bucket...')
tmp_valid_files = []
current_date = start_date
while current_date <= end_date:
    date = current_date.strftime("%Y%m%d")
    day_of_year = current_date.strftime('%j')

    year  = int(date[0:4])
    month = int(date[4:6])
    day   = int(date[6:8])

    print('Checking', date, 'on S3 bucket...')
    # Check if there is data on the server for the requested day.
    # If not, skip this day.
    number_of_obj = 0
    # Need to loop through the hours for the NODD S3 buckets
    for h in range(0,24):
        hstr = '{:02d}'.format(h)
        this_prefix = wl_folders[wavelength]+'/'+'{:04d}'.format(year)+'/'+day_of_year+'/'+hstr+'/'

        for obj in suvi_bucket.objects.filter(Prefix=this_prefix).all():
            if obj.key.endswith(file_ending):
                number_of_obj += 1
                tmp_valid_files.append(obj.key)

    if number_of_obj == 0:
        warn_msg = f'No files found on S3 bucket for {date}. Skipping...'
        print(warn_msg)
        current_date = current_date + datetime.timedelta(days=1)
        continue

    current_date = current_date + datetime.timedelta(days=1)

# Get the start and end times from the filenames on the S3 bucket
valid_files = []
if len(tmp_valid_files) == 0:
    err_msg = 'No valid files could be selected.'
    raise Exception(err_msg)
else:
    print('Found', str(len(tmp_valid_files)), 'valid files on the S3.')
    tmp_start_times = []
    tmp_end_times = []
    for tmp_file in tmp_valid_files:
        file_base = os.path.basename(tmp_file)
        file_base_split = file_base.split('_')
        stime = file_base_split[3][1:]
        etime = file_base_split[4][1:]
        tmp_start_time = stime[0:4]+':'+stime[4:7]+':'+stime[7:9]+':'+stime[9:11]+':'+stime[11:13]
        tmp_end_time   = etime[0:4]+':'+etime[4:7]+':'+etime[7:9]+':'+etime[9:11]+':'+etime[11:13]
        # Exclude files outside of the requested period if requested:
        if enforce_requested_period:
            if ((datetime.datetime.strptime(tmp_start_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc)-start_date).total_seconds() >= 0) and \
               ((datetime.datetime.strptime(tmp_end_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc)-end_date).total_seconds() <= 0):
                valid_files.append(tmp_file)
                tmp_start_times.append(datetime.datetime.strptime(tmp_start_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc))
                tmp_end_times.append(datetime.datetime.strptime(tmp_end_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc))
        else:
            valid_files.append(tmp_file)
            tmp_start_times.append(datetime.datetime.strptime(tmp_start_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc))
            tmp_end_times.append(datetime.datetime.strptime(tmp_end_time, "%Y:%j:%H:%M:%S").replace(tzinfo=datetime.timezone.utc))

    if len(valid_files) == 0:
        err_msg = 'No valid files found on the NODD S3 buckets.'
        raise Exception(err_msg)

    # Make lists of start times, end times, and exposure times
    start_time = []
    end_time = []
    exposure_time = []
    for counter, tmp_start_time in enumerate(tmp_start_times):
        start_time.append(tmp_start_time)
        end_time.append(tmp_end_times[counter])
        exposure_time.append((tmp_end_times[counter]-tmp_start_time).total_seconds())

    exposure_time = np.array(exposure_time)

    # Select long or short exposures, depending on what was requested
    if short_exp:
        chosen_exposures = np.where(np.around(exposure_time) == 0.)
    else:
        chosen_exposures = np.where(np.around(exposure_time) == 1.)

    chosen_exposure_files = list(np.array(valid_files)[chosen_exposures])

# Make the reference timestamps at the requested cadence
reference_timestamps = []
current_timestamp = start_date
while current_timestamp <= end_date:
    reference_timestamps.append(current_timestamp)
    current_timestamp = current_timestamp + datetime.timedelta(minutes=cadence)

selected_files = []
for ttime in reference_timestamps:
    delta_t = np.zeros(len(chosen_exposures[0]))
    for i in range(0,len(chosen_exposures[0])):
        delta_t[i] = (ttime-start_time[chosen_exposures[0][i]]).total_seconds()
    # Only accept files within the given tolerance:
    if (np.abs(delta_t).min()/60.) <= tolerance:
        which_file = np.abs(delta_t).argmin()
        # Avoid adding the same file multiple times
        if chosen_exposure_files[which_file] not in selected_files:
            selected_files.append(chosen_exposure_files[which_file])

if len(selected_files) == 0:
    err_msg = 'No valid files could be selected.'
    raise Exception(err_msg)
else:
    if no_sub_dirs:
        store_loc = outpath
    else:
        store_loc = outpath+satellite_name_long+'/'+date[0:4]+'/'+date[4:6]+'/'+'/'+date[6:8]+'/'+str(wavelength)+'/'

    if not dryrun:
        # Create folders for downloaded data if they do not exist
        if not os.path.exists(store_loc):
            os.makedirs(store_loc)

    for sf in selected_files:
        print('Downloading', os.path.basename(sf))
        if not dryrun:
            suvi_bucket.download_file(sf, store_loc+os.path.basename(sf))