import numpy as np
import tifffile
import astropy
from astropy.io import fits
from read_roi import read_roi_file
import logging
import os
import types
import inspect
from neutompy.misc.uitools import get_filename_gui, get_image_gui, get_folder_gui
from tqdm import tqdm

__author__  = "Davide Micieli"
__all__     = ['read_tiff',
			 'read_fits',
			 'read_fits_stack',
			 'read_tiff_stack',
			 'get_rect_coordinates_from_roi',
			 'read_image',
			 'read_stack_from_list',
			 'read_dataset',
			 'read_image_stack',
			 'write_fits',
			 'write_tiff',
			 'write_tiff_stack',
			 'write_fits_stack',
			 'get_rect_coordinates_from_roi', 'get_filename_pattern']

#logging.basicConfig(level=logging.WARNING)
logs = logging.getLogger(__name__)


known_ext = ['.tif', '.tiff', '.fits']

# to ensure compatibility with astropy 1.X, clobber option is deprecated and replaced with overwrite
if 'clobber' in inspect.getfullargspec(astropy.io.fits.writeto).args:
	arg_overwrite = 'clobber'
else:
	arg_overwrite = 'overwrite'


def get_file_extension(fname):
	"""
	This function returns the extension of a given file name or file path
	"""
	return os.path.splitext(fname)[1]


def get_rect_coordinates_from_roi (fname):
	"""
		This function returns the coordinates from a rectangular region of interest
		defined in a .roi file generated by ImageJ / Fiji.

		N.B.: indexing starts from 0.
		Incremet rowmax and col_max  + 1 to crop an image, for example:
		crop = img[rowmin:(rowmax+1),  colmin:(colmax+1)]

		Parameters
		----------
		fname : str
			String defining the path or the name of the Image ROI file.

		Returns
		-------
		rowmin : int
			The minimum row coordinate.

		rowmax : int
			The maximum row coordinate.

		colmin : int
			The minimum column coordinate.

		colmax : int
			The maximum column coordinate.

	"""

	# file name check
	if(type(fname) is not str):
		logs.error('File name must be a string.')
		raise TypeError('File name must be a string.')
	else:
		if not(fname.endswith('.roi')):
			logs.warning('File extension shold be .roi')

	# read roi file
	roi = read_roi_file(fname)


	# convert keys to list
	l = list(roi.keys())

	# extract coordinates and length
	height = roi[l[0]]['height']
	width  = roi[l[0]]['width']
	top    = roi[l[0]]['top']
	left	 = roi[l[0]]['left']

	rowmin = int(top)
	rowmax = int(top  + height - 1)
	colmin = int(left)
	colmax = int(left + width - 1)

	# log infos
	logs.info('ROI file succesfully read: %s', fname)

	return rowmin, rowmax, colmin, colmax


def read_tiff(fname, croi=None, froi=None):
	"""
	This function reads a 2D TIFF image.

	Parameters
	----------
	fname : str
		String defining the file name or the file path

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI file name or file path.

	Returns
	-------
	out : ndarray
		The 2D image.
		If the file specified doesn't exist, the function returns False.

	"""
	# file name check
	if(type(fname) is not str):
		logs.error('File name must be a string.')
		raise TypeError('File name must be a string.')

	if ( not(fname.endswith('.tif')) and  not(fname.endswith('.tiff')) ):
		logs.warning('File extension shold be .tif or .tiff')

	# ROI check
	if(croi and froi):
		raise ValueError('Please set the coordinates of the region of interest OR the ImageJ file .roi')

	if(croi):
		if not(len(croi)==2 and len(croi[0])==3  and len(croi[1])==3  ):
			raise ValueError('Tuple or list size not valid. Please set the coordinates of the ROI in the this manner:   croi = ( (row_start, row_end, row_step),  (col_start, col_end, col_step) ) ')


	# get the absolute path of the file
	fname = os.path.abspath(fname)

	try:
		out = tifffile.imread(fname) #, out='memmap')
		logs.info('Array imported from file: %s', fname)
		logs.debug('Array shape and type: %s, %s', out.shape, out.dtype)
	except IOError:
		logs.error('No such file: %s', fname)
		return False

	if (len(out.shape) != 2):
		raise ValueError('Only 2D TIFF images are supported in read_tiff')

	if(croi):
		rstart, rend, rstep = croi[0]
		cstart, cend, cstep = croi[1]
		if (rend is not None):
			rend = rend + 1
		if (cend is not None):
			cend = cend + 1
		out  = out[rstart:rend:rstep, cstart:cend:cstep]

		logs.info('Read roi coordinate from list / tuple')
		logs.debug('Cropped array shape and type: %s, %s', out.shape, out.dtype)
		logs.debug('rstart =  %s, rend = %s, rstep = %s', rstart, rend, rstep)
		logs.debug('cstart =  %s, cend = %s, cstep = %s', cstart, cend, cstep)

	if(froi):
		rstart, rend, cstart, cend = get_rect_coordinates_from_roi(froi)
		out = out[rstart:rend + 1, cstart:cend + 1]

		logs.info('Read ROI coordinate from file : %s' , froi)
		logs.debug('Cropped array shape and type: %s, %s', out.shape, out.dtype)
		logs.debug('rstart =  %s, rend = %s, rstep = %s', rstart, rend)
		logs.debug('cstart =  %s, cend = %s, cstep = %s', cstart, cend)


	return out

def read_fits(fname, croi=None, froi=None):
	"""
	This function reads a 2D FITS image.

	Parameters
	----------
	fname : str
		String defining the file name or the file path

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	out : ndarray
		The 2D image.
		If the file specified doesn't exist, the function returns False.

	"""
	# file name check
	if(type(fname) is not str):
		logs.error('File name must be a string.')
		raise TypeError('File name must be a string.')

	if ( not(fname.endswith('.fits')) ):
		logs.warning('File extension shold be .fits')

	# ROI check
	if(croi and froi):
		raise ValueError('Please set the coordinates of the region of interest OR the ImageJ file .roi')

	if(croi):
		if not(len(croi)==2 and len(croi[0])==3  and len(croi[1])==3  ):
			raise ValueError('Tuple or list size not valid. Please set the coordinates of the ROI in the this manner:   croi = ( (row_start, row_end, row_step),  (col_start, col_end, col_step) ) ')




	# get the absolute path of the file
	fname = os.path.abspath(fname)

	try:
		f = fits.open(fname)
		out = f[0].data[::-1, :].copy() # flip up-down
		f.close()
		logs.info('Array imported from file: %s', fname)
		logs.debug('Array shape and type: %s, %s', out.shape, out.dtype)
	except IOError:
		logs.error('No such file: %s', fname)
		return False

	if (len(out.shape) != 2):
		raise ValueError('Only 2D FITS images are supported in read_tiff')

	if(croi):
		rstart, rend, rstep = croi[0]
		cstart, cend, cstep = croi[1]
		if (rend is not None):
			rend = rend + 1
		if (cend is not None):
			cend = cend + 1
		out  = out[rstart:rend:rstep, cstart:cend:cstep]

		logs.info('Read roi coordinate from list / tuple')
		logs.debug('Cropped array shape and type: %s, %s', out.shape, out.dtype)
		logs.debug('rstart =  %s, rend = %s, rstep = %s', rstart, rend, rstep)
		logs.debug('cstart =  %s, cend = %s, cstep = %s', cstart, cend, cstep)

	if(froi):
		rstart, rend, cstart, cend = get_rect_coordinates_from_roi(froi)
		out = out[rstart:rend + 1, cstart:cend + 1]

		logs.info('Read ROI coordinate from file : %s' , froi)
		logs.debug('Cropped array shape and type: %s, %s', out.shape, out.dtype)
		logs.debug('rstart =  %s, rend = %s, rstep = %s', rstart, rend)
		logs.debug('cstart =  %s, cend = %s, cstep = %s', cstart, cend)

	return out



def read_image(fname, croi=None, froi=None):
	"""
	This function reads a 2D TIF or FITS image.

	Parameters
	----------
	fname : str
		String defining the file name or the file path

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	out : ndarray
		The 2D image.
		If the file specified doesn't exist, the function returns False.

	"""

	# check if file exists
	if not(os.path.isfile(fname)):
		raise ValueError('No such file %s' %fname)

	# read .fits or .tif from the file extension
	if(fname.endswith('.fits') or fname.endswith('.FITS')):
		funz = read_fits
	elif(fname.endswith('.tif') or fname.endswith('.tiff') or fname.endswith('.TIF') or fname.endswith('.TIFF')):
		funz = read_tiff
	else:
		raise ValueError('File extension not valid. File : %s' % fname)

	array = funz(fname, croi, froi)

	return array



def get_filename_pattern(fname):
	# get file name extension
	ext = get_file_extension(fname)

	# get prefix_XXXX
	aus = fname[:-len(ext)]

	# find how many X, remove XXXX and find the prefix name
	p = 0
	while(aus[-1-p].isdigit()):
		p = p + 1
	prefix_full = aus[:-p]

	folder = os.path.dirname(fname)
	prefix = os.path.basename(prefix_full)

	numbering = '#'*p

	return prefix_full, numbering, ext


def read_stack_from_list(flist, slices=[], croi=None, froi=None):
	"""
	Read stack of images from file list.

	Parameters
	----------
	flist : list of str
		List of the file names or file paths to read in the defined order. File with unknown image extension are skipped automatically.

	slices : list of int, optional
		List of the indexes of flist to read. If it is [] then all image files are read.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	array : ndarray
		The 3D stack of images

	"""

	# if slices == [], initialize all indexes
	if not slices:
		slices = [ i for i in range(0, len(flist)) ]

	# list of index to read
	toread = [True]*len(slices)

	# skip files with unknown extension
	for e, s in enumerate(slices):
		if not (get_file_extension(flist[s]) in known_ext):
			toread[e] = False
			logs.warning('Unknown file extension. Skipped file: %s', flist[s])


	fpat = get_filename_pattern(flist[0])  # e.g. /folder1/folder2/img_####.tiff

	print('> Reading stack of images:', ''.join(fpat))
	array = np.array([read_image(flist[s], croi, froi) for e, s in enumerate(tqdm(slices, unit=' images')) if toread[e]])
	print('')

	return array


def read_tiff_stack(fname, slices=[], croi=None, froi=None):
	"""
	Read stack of tiff images.

	Parameters
	----------
	fname :  str
		One of the file names of the tiff stack. File with unknown image extension are skipped automatically. If fname is '' or a folder path, then a dialog box is opened to select one of the file of the stack to read. The initial directory is C:\ (in windows) or / (in UNIX) if fname=='', otherwise is the folder path assigned to fname.

	slices : list of int, optional
		List of the element indexes to read. If it is [] then all tif image files are read.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	array : ndarray
		The 3D stack of images

	"""

	if(os.path.isdir(fname) or fname==''):
		fname = get_filename_gui(initialdir=fname, message='Select a file of the image stack...',
					ext=('TIFF Images', '*.tif *.tiff') )

	# check if file exists
	if not(os.path.isfile(fname)):
		raise ValueError('No such file %s' %fname)

	# get file name extension
	ext = get_file_extension(fname)

	# get prefix_XXXX
	aus = fname[:-len(ext)]

	# find how many X, remove XXXX and find the prefix name
	p = 0
	while(aus[-1-p].isdigit()):
		p = p + 1
	prefix_full = aus[:-p]

	folder = os.path.dirname(fname)
	prefix = os.path.basename(prefix_full)

	flist = [os.path.normpath(os.path.join(folder, f)) for f in sorted(os.listdir(folder)) if (f.startswith(prefix) and f[-p-len(ext):-len(ext)].isdigit() and (get_file_extension(f) in ['.tif', '.tiff']) ) ]

	out = read_stack_from_list(flist, slices, croi, froi)

	return out


def read_fits_stack(fname, slices=[], croi=None, froi=None):
	"""
	Read stack of fits images.

	Parameters
	----------
	fname :  str
		One of the file names of the fits stack. File with unknown image extension are skipped automatically. If fname is '' or a folder path, then a dialog box is opened to select one of the file of the stack to read. The initial directory is C:\ (in windows) or / (in UNIX) if fname=='', otherwise is the folder path assigned to fname.

	slices : list of int, optional
		List of the element indexes to read. If it is [] then all fits image files are read.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	array : ndarray
		The 3D stack of images

	"""

	if(os.path.isdir(fname) or fname==''):
		fname = get_filename_gui(initialdir=fname, message='Select a file of the image stack...', ext=('FITS Images', '*.fits') )

	# check if file exists
	if not(os.path.isfile(fname)):
		raise ValueError('No such file %s' %fname)

	# get file name extension
	ext = get_file_extension(fname)

	# get prefix_XXXX
	aus = fname[:-len(ext)]

	# find how many X, remove XXXX and find the prefix name
	p = 0
	while(aus[-1-p].isdigit()):
		p = p + 1
	prefix_full = aus[:-p]

	folder = os.path.dirname(fname)
	prefix = os.path.basename(prefix_full)

	flist = [os.path.normpath(os.path.join(folder, f)) for f in sorted(os.listdir(folder)) if (f.startswith(prefix) and f[-p-len(ext):-len(ext)].isdigit() and (get_file_extension(f) in ['.fits']) ) ]

	out = read_stack_from_list(flist, slices, croi, froi)

	return out

def read_image_stack(fname, slices=[], croi=None, froi=None):
	"""
	Read stack of TIFF or FITS images. This function recognize the file type from
	the file extension.

	Parameters
	----------
	fname :  str
		One of the file names of the tiff stack. File with unknown image extension are
		skipped automatically.

	slices : list of int, optional
		List of the element indexes to read. If it is [] then all tif image files are read.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI filename or file path.

	Returns
	-------
	array : ndarray
		The 3D stack of images

	"""
	if(fname.endswith('.fits') or fname.endswith('.FITS')):
		funz = read_fits_stack
	elif(fname.endswith('.tif') or fname.endswith('.tiff') or fname.endswith('.TIFF') or fname.endswith('.TIF')):
		funz = read_tiff_stack
	else:
		raise ValueError('File extension not valid. File : %s' % fname)

	arr = funz(fname, slices, croi, froi)

	return arr


def read_dataset(proj_180=True, croi=None, froi=None):
	"""
	This function reads a dataset which contains dark-field, flat-field, projection
	images and optionally the projection at 180 degree. The user selects the main folder
	and the files from a dialog box. A two-dimensional region of interest of each
	image can be read specifying the coordinates or an ImageJ .roi file.

	Parameters
	----------
	proj_180 : bool, optional
		If True the user must select the projection at 180 degree separately from
		the stack of projections.

	croi : tuple, optional
		Tuple defining the ROI indexes for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	froi : str, optional
		String defining the ImageJ ROI file name or file path.

	Returns
	-------
	proj : 3d array
		The array containing the stack of projections.

	dark : 3d array
		The array containing the stack of dark-field images.

	flat : 3d array
		The array containing the stack of flat-field images.

	proj180 : 2d array
		Only if proj_180 is True, the 2D array representing the projection at 180 degree
		is returned separately.
	"""
	# select the folder that contains dark, flat, and projection images using the GUI
	data_folder = get_folder_gui('', message = 'Select the main folder of the dataset...')

	proj_dir = get_image_gui(data_folder, message = 'Select a raw projection...')
	dark_dir = get_image_gui(data_folder, message = 'Select a dark image...')
	flat_dir = get_image_gui(data_folder, message = 'Select a flat image...')
	if proj_180:
		proj_pi_fname =  get_image_gui(data_folder, message = 'Select raw projection at 180°...')

	# load the dataset
	proj = read_image_stack(proj_dir, croi=croi, froi=froi)
	dark = read_image_stack(dark_dir, croi=croi, froi=froi)
	flat = read_image_stack(flat_dir, croi=croi, froi=froi)
	if proj_180:
		print('> Reading projection at 180 degree:',proj_pi_fname)
		proj180 = read_image(proj_pi_fname, croi=croi, froi=froi)

	if proj_180:
		return proj, dark, flat, proj180
	else:
		return proj, dark, flat


def write_tiff(fname, img, overwrite=False):
	"""
	This function write to the disk an array as tiff image.

	Parameters
	----------
	fname : str
		String defining the file name or file path of the image to save.
		If the extension is not specified, it is automatically appended to fname.

	img : ndarray
		The array to save as image.

	overwrite: bool, optional
		If ``True``, overwrites the output file if it exists. Raises an
		``IOError`` if ``False`` and the output file exists. Default is ``False``.
	"""

	# add extension .tif if not specified in fname
	if not ( fname.endswith('.tif') or fname.endswith('.tiff') ):
		fname = fname + '.tiff'

	if not overwrite:
		if(os.path.isfile(fname)):
			raise IOError("File already exists : %s" %fname)

	if overwrite:
		if(os.path.isfile(fname)):
			logs.info("File to overwrite: %s", fname)

	tifffile.imsave(fname, img)
	logs.info('File saved to disk: %s', fname)
	logs.debug('Image shape %s and type %s', img.shape,  img.dtype)


def write_fits(fname, img, overwrite=False):
	"""
	This function write to the disk an array as fits image.

	Parameters
	----------
	fname : str
		String defining the file name or file path of the image to save.
		If the extension is not specified, it is automatically appended to fname.

	img : ndarray
		The array to save as image.

	overwrite: bool, optional
		If ``True``, overwrites the output file if it exists. Raises an ``IOError``
		if ``False`` and the output file exists. Default is ``False``.
	"""

	# add extension .tif if not specified in fname
	if not ( fname.endswith('.fits') ):
		fname = fname + '.fits'

	if overwrite:
		if(os.path.isfile(fname)):
			logs.info("File to overwrite: %s", fname)

	hdu = fits.PrimaryHDU()
	hdu.data = img[::-1, :]
	hdu.writeto(fname, **{arg_overwrite : overwrite})

	logs.info('File saved to disk: %s', fname)
	logs.debug('Image shape %s and type %s', img.shape,  img.dtype)


def write_tiff_stack(fname, data, axis=0, start=0, croi=None, digit=4, dtype=None, overwrite=False):
	"""
	This function writes a 3D array to a stack of 2D TIFF images.

	Parameters
	----------

	fname: str
		String defining the prefix of the file name and the containing folder.

	data : ndarray
		The 3D stack to write.

	axis : int, optional
		The axis along which the stacking is performed.

	start : int, optional
		Index used for saving the first image.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	digit : int, optional
		Number of digits used for the numbering of the images.

	dtype : dtype, optional
		Data type of the images to save.

	overwrite: bool, optional
		If ``True``, overwrites the output file if it exists. Raises an ``IOError``
		if ``False`` and the output file exists. Default is ``False``.

	"""
	## check dimension data ==3
	if (data.ndim != 3):
		raise ValueError('Array must have 3 dimensions.')

	folder = os.path.dirname(fname)

	if folder == '':
		raise ValueError('File path not valid. Please specify the file path giving at\
				least one folder, e.g.: folder/file.tiff ')
	else:
		if not (os.path.isdir(folder)):
			raise ValueError('Folder not exist. %s' %folder)

	if not dtype:
		dtype = data.dtype


	if(axis!=0):
		data  = np.swapaxes(data, 0, axis)
		if(croi):
			newcroi_list       = [[None]*3]*3
			newcroi_list[axis] = croi[0]
			newcroi_list[0]    = croi[axis]
			newcroi_list[3-axis] = croi[3-axis]
			newcroi              = tuple(newcroi_list)

	if not croi:
		slices, rows, cols = data.shape
		newcroi = ((0, slices, 1), (0, rows, 1), (0, cols, 1))
	else:
		if(axis ==0):
			newcroi = croi

	rmin, rmax, rstep = newcroi[1]
	cmin, cmax, cstep = newcroi[2]
	smin, smax, sstep = newcroi[0]

	auslist = list(range(0, data.shape[0], 1))

	subset = auslist[smin:smax:sstep]

	logs.debug('Shape of the re-stacked array %s',  newcroi)

	print('> Saving stack of images:', fname + '_' + '#'*digit + '.tiff')
	for iz, s in enumerate(tqdm(subset, unit=' images'), start):
		outfile = fname + '_' + str(iz).zfill(digit) + '.tiff'
		write_tiff(outfile, data[s, rmin:rmax:rstep, cmin:cmax:cstep].astype(dtype), overwrite)


def write_fits_stack(fname, data, axis=0, start=0, croi=None, digit=4, dtype=None, overwrite=False):
	"""
	This function writes a 3D array to a stack of 2D FITS images.

	Parameters
	----------

	fname: str
		String defining the prefix of the file name and the containing folder.

	data : ndarray
		The 3D stack to write.

	axis : int, optional
		The axis along which the stacking is performed.

	start : int, optional
		Index used for saving the first image.

	croi : tuple, optional
		Tuple defining the indexes range for each axis. It must be follow this notation:
		( (row_start, row_end, row_step),  (col_start, col_end, col_step) )

	digit : int, optional
		Number of digits used for the numbering of the images.

	dtype : dtype, optional
		Data type of the images to save.

	overwrite: bool, optional
		If ``True``, overwrites the output file if it exists. Raises an ``IOError``
		if ``False`` and the output file exists. Default is ``False``.

	"""
	## check dimension data ==3
	if (data.ndim != 3):
		raise ValueError('Array must have 3 dimensions.')

	folder = os.path.dirname(fname)

	if folder == '':
		raise ValueError('File path not valid. Please specify the file path giving\
					at least one folder, e.g.: folder/file.fits ')
	else:
		if not (os.path.isdir(folder)):
			raise ValueError('Folder not exist. %s' %folder)

	if not dtype:
		dtype = data.dtype


	if(axis!=0):
		data  = np.swapaxes(data, 0, axis)
		if(croi):
			newcroi_list       = [[None]*3]*3
			newcroi_list[axis] = croi[0]
			newcroi_list[0]    = croi[axis]
			newcroi_list[3-axis] = croi[3-axis]
			newcroi              = tuple(newcroi_list)

	if not croi:
		slices, rows, cols = data.shape
		newcroi = ((0, slices, 1), (0, rows, 1), (0, cols, 1))
	else:
		if(axis ==0):
			newcroi = croi

	rmin, rmax, rstep = newcroi[1]
	cmin, cmax, cstep = newcroi[2]
	smin, smax, sstep = newcroi[0]

	auslist = list(range(0, data.shape[0], 1))

	subset = auslist[smin:smax:sstep]

	logs.debug('Shape of the re-stacked array %s',  newcroi)

	print('> Saving stack of images:', fname + '_' + '#'*digit + '.fits')
	for iz, s in enumerate(tqdm(subset, unit=' images'), start):
		outfile = fname + '_' + str(iz).zfill(digit) + '.fits'
		write_fits(outfile, data[s, rmin:rmax:rstep, cmin:cmax:cstep].astype(dtype), overwrite)
