#################################################################################
from pathlib import Path
import os
import sys
from astropy.stats import mad_std
from astropy.io import fits
from astropy.utils.data import get_pkg_data_filename
import glob
from astropy import stats
import matplotlib.pyplot as plt
import numpy as np
from datetime import date
today = date.today()
from astropy.wcs import WCS
import warnings
warnings.filterwarnings('ignore')
from matplotlib.backends.backend_pdf import PdfPages
import shutil
#################################################################################

# Read in day number from command line:  

n = len(sys.argv)
if n == 0 or n == 1 or n > 3:
    print('Usage:  python extract_subimages.ipynb ### $$$' )
    print('   where ### is a 3-digit identifier for the day number')
    print('   and $$$ is an (optional) 3-digit identifier for the observer code')
    # end program:
    sys.exit(0)

if n == 2:  
    daynum = sys.argv[1]
    obscode = '*'
if n == 3:  
    daynum = sys.argv[1]
    obscode = sys.argv[2]

#################################################################################
# Use glob to create list:
current_dir = os.getcwd()
if obscode == '*':
    image_dir = '/mnt/imagesbucket/*/*/2023/'+str(daynum)+'/'
if obscode != '*':
    image_dir = '/mnt/imagesbucket/*/'+str(obscode)+'/2023/'+str(daynum)+'/'
out_dir = current_dir+'/'
os.makedirs(out_dir,exist_ok = True)
image_list = sorted(glob.glob(image_dir+'*.fts',recursive=True))
N_images = len(image_list)
print('-----------------------------------------------------------')
if obscode == '*':
    print('    --> '+str(N_images)+' images found for day '+str(daynum)+' from all observers')
if obscode != '*':
    print('    --> '+str(N_images)+' images found for day '+str(daynum)+' from observer '+str(obscode))
print('-----------------------------------------------------------')
#################################################################################

# extract central subimage from each image:
size = 256 # number of pixels each quadrant of subimage
for i in range(N_images):
    image = image_list[i]
    imagename = str.split(image,'/')[7]
    imageroot = str.split(imagename,'.')[0]
    #print(imageroot)
    hdu = fits.open(image)
    data = hdu[0].data
    header = hdu[0].header
    hdu.close()
    # get image dimensions:
    nx = data.shape[0]
    ny = data.shape[1]
    # get central pixel coordinates:
    x0 = nx/2
    y0 = ny/2
    # extract subimage:
    subimage = data[int(x0-size):int(x0+size),int(y0-size):int(y0+size)]
    # write subimage to file:
    subimage_file = out_dir+str(imageroot)+'.trim.fits'
    hdu = fits.PrimaryHDU(subimage)
    hdu.writeto(subimage_file,overwrite=True)
    #hdu.close()
    print('      '+str(imagename)+' ('+str(nx)+'x'+str(ny)+') --->> '+str(imageroot)+'.trim.fits ('+str(size*2)+'x'+str(size*2)+')')