import numpy as np

def read_fits(filename):
	'''Read an array from a fits file.

	Parameters
	----------
	filename : string
		The filename of the file to read. This can include a path.

	Returns
	-------
	ndarray
		The ndarray read from the fits file.
	'''
	from astropy.io import fits
	return fits.getdata(filename).copy()

def write_fits(data, filename, shape=None, overwrite=True):
	'''Write the data to a fits-file.

	If the data is a Field with SeparatedCoords, the shaped field will be written to the
	fits-file, to allow for easier viewing with external tools.

	Parameters
	----------
	data : ndarray or Field
		The ndarray or Field to write to the fits file.
	filename : string
		The filename of the newly created file. This can include a path.
	shape : ndarray or None
		The shape to which to reshape the data. If this is given, it will override a potential
		shape from the grid accompaning the field.
	overwrite : boolean
		Whether to overwrite the fits-file if it already exists.
	'''
	from astropy.io import fits

	hdu = None

	if shape is not None:
		hdu = fits.PrimaryHDU(data.reshape(shape))
	elif hasattr(data, 'grid'):
		if data.grid.is_separated:
			hdu = fits.PrimaryHDU(data.shaped)

	if hdu is None:
		hdu = fits.PrimaryHDU(data)

	hdu.writeto(filename, overwrite=True)

def _guess_file_format(filename):
	'''Guess the file format from the filename.

	Parameters
	----------
	filename : string
		The filename of the file.

	Returns
	-------
	string
		The file type. This is one of ["asdf", "fits"].
	'''
	if filename.endswith('asdf'):
		return 'asdf'
	elif filename.endswith('fits') or filename.endswith('fits.gz'):
		return 'fits'
	elif filename.endswith('pkl') or filename.endswith('pickle'):
		return 'pickle'
	else:
		return None

def _make_metadata(file_type):
	'''Make a metadata section for file output.

	This contains the version of HCIPy used to write the file, the current time and
	date and the file type.

	Parameters
	----------
	file_type : string
		The identifier string for what is going to be stored in the file.

	Returns
	-------
	dictionary
		The metadata.
	'''
	from ..version import get_version
	import datetime

	tree = {
		'meta': {
			'author': 'HCIPy %s' % get_version(),
			'date_utc': datetime.datetime.utcnow().isoformat(),
			'file_type': file_type
		}
	}

	return tree

def read_grid(filename, fmt=None):
	'''Read a grid from a file.

	Parameters
	----------
	filename : string
		The path of the file you want to read the grid from.
	fmt : string
		The file format. If it is not given, the file format will be guessed from the file extension.

	Returns
	-------
	Grid
		The read grid.

	Raises
	------
	ValueError
		If the file format could not be guessed from the file extension.
	NotImplementedError
		If the file format was not yet implemented.
	'''
	from ..field import Grid

	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	if fmt in ['asdf', 'fits']:
		import asdf

		f = asdf.open(filename)
		grid = Grid.from_dict(f.tree['grid'])
		f.close()

		return grid
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'rb') as f:
			return pickle.load(f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)

def write_grid(grid, filename, fmt=None, overwrite=True):
	'''Write a grid to a file.

	Parameters
	----------
	grid : Grid
		The grid to write to the file.
	filename : string
		The path of the file you want to write to.
	fmt : string
		The file format. If it is not given, the file format will be guessed based on the file extension.
	overwrite : boolean
		Whether to overwrite the file, if it already exists.

	Raises
	------
	NotImplementedError
		If the file format was not yet implemented.
	'''
	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	tree = _make_metadata('grid')
	tree['grid'] = grid.to_dict()

	if fmt == 'asdf':
		import asdf

		target = asdf.AsdfFile(tree)
		target.write_to(filename, all_array_compression='zlib')
	elif fmt == 'fits':
		from astropy.io import fits
		import asdf

		hdulist = fits.HDUList()

		ff = asdf.fits_embed.AsdfInFits(hdulist, tree)
		ff.write_to(filename, all_array_compression='zlib', overwrite=overwrite)
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'wb') as f:
			pickle.dump(grid, f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)

def read_field(filename, fmt=None):
	'''Read a field from a file.

	Parameters
	----------
	filename : string
		The path of the file you want to read the field from.
	fmt : string
		The file format. If it is not given, the file format will be guessed from the file extension.

	Returns
	-------
	Field
		The read field.

	Raises
	------
	ValueError
		If the file format could not be guessed from the file extension.
	NotImplementedError
		If the file format was not yet implemented.
	'''
	from ..field import Field

	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	if fmt == 'asdf':
		import asdf

		f = asdf.open(filename)
		field = Field.from_dict(f.tree['field'])
		f.close()

		return field
	elif fmt == 'fits':
		import asdf
		from ..field import Grid

		f = asdf.open(filename)
		tree = f.tree['field']

		if 'grid' in tree:
			grid = Grid.from_dict(tree['grid'])
			new_shape = np.concatenate((tree['values'].shape[:-grid.ndim], [grid.size])).astype('int')
			tree['values'] = tree['values'].reshape(new_shape)

		field = Field.from_dict(tree).reshape(new_shape)
		f.close()

		return field
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'rb') as f:
			return pickle.load(f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)

def write_field(field, filename, fmt=None, overwrite=True):
	'''Write a field to a file.

	Parameters
	----------
	field : Field
		The field to write to the file.
	filename : string
		The path of the file you want to write to.
	fmt : string
		The file format. If it is not given, the file format will be guessed based on the file extension.
	overwrite : boolean
		Whether to overwrite the file, if it already exists.

	Raises
	------
	NotImplementedError
		If the file format was not yet implemented.
	'''
	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	tree = _make_metadata('field')
	tree['field'] = field.to_dict()

	if fmt == 'asdf':
		import asdf

		target = asdf.AsdfFile(tree)
		target.write_to(filename, all_array_compression='zlib')
	elif fmt == 'fits':
		from astropy.io import fits
		import asdf

		hdulist = fits.HDUList()

		if field.grid.is_separated:
			hdulist.append(fits.ImageHDU(np.ascontiguousarray(field.shaped)))
			tree['field']['values'] = hdulist[0].data

		if field.grid.is_regular:
			from astropy import wcs

			w = wcs.WCS(naxis=field.grid.ndim)
			w.wcs.crpix = np.ones(field.grid.ndim)
			w.wcs.cdelt = field.grid.delta
			w.wcs.crval = field.grid.zero
			w.wcs.ctype = ['X', 'Y', 'Z', 'W'][:field.grid.ndim]

			hdulist[0].header.update(w.to_header())

		ff = asdf.fits_embed.AsdfInFits(hdulist, tree)
		ff.write_to(filename, all_array_compression='zlib', overwrite=overwrite)
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'wb') as f:
			pickle.dump(field, f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)

def read_mode_basis(filename, fmt=None):
	'''Read a mode basis from a file.

	Parameters
	----------
	filename : string
		The path of the file you want to read the mode basis from.
	fmt : string
		The file format. If it is not given, the file format will be guessed from the file extension.

	Returns
	-------
	ModeBasis
		The read mode basis.

	Raises
	------
	ValueError
		If the file format could not be guessed from the file extension.
	NotImplementedError
		If the file format was not yet implemented.
	'''
	from ..mode_basis import ModeBasis

	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	if fmt == 'asdf':
		import asdf

		f = asdf.open(filename)
		mode_basis = ModeBasis.from_dict(f.tree['mode_basis'])
		f.close()

		return mode_basis
	elif fmt == 'fits':
		import asdf
		from ..field import Grid

		f = asdf.open(filename)
		tree = f.tree['mode_basis']

		if 'modes' in tree:
			grid = Grid.from_dict(tree['grid'])
			modes = tree['modes']

			old_shape = np.concatenate((modes.shape[:-grid.ndim], [grid.size]))
			tree['transformation_matrix'] = modes.reshape(old_shape).T
			del tree['modes']

		mode_basis = ModeBasis.from_dict(tree)
		f.close()

		return mode_basis
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'rb') as f:
			return pickle.load(f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)

def write_mode_basis(mode_basis, filename, fmt=None, overwrite=True):
	'''Write a mode basis to a file.

	Parameters
	----------
	mode_basis : ModeBasis
		The mode basis to write to the file.
	filename : string
		The path of the file you want to write to.
	fmt : string
		The file format. If it is not given, the file format will be guessed based on the file extension.
	overwrite : boolean
		Whether to overwrite the file, if it already exists.

	Raises
	------
	NotImplementedError
		If the file format was not yet implemented.
	'''
	if fmt is None:
		fmt = _guess_file_format(filename)

		if fmt is None:
			raise ValueError('Format not given and could not be guessed based on the file extension.')

	tree = _make_metadata('mode_basis')
	tree['mode_basis'] = mode_basis.to_dict()

	if fmt == 'asdf':
		import asdf

		target = asdf.AsdfFile(tree)
		target.write_to(filename, all_array_compression='zlib')
	elif fmt == 'fits':
		from astropy.io import fits
		import asdf
		from astropy import wcs

		hdulist = fits.HDUList()

		if mode_basis.grid and mode_basis.grid.is_separated:
			new_shape = np.concatenate(([-1], mode_basis.grid.shape))
			modes = np.ascontiguousarray(mode_basis.to_dense().transformation_matrix.T.reshape(new_shape))

			hdulist.append(fits.ImageHDU(np.asarray(modes)))

			w = wcs.WCS(naxis=mode_basis.grid.ndim)
			w.wcs.crpix = np.ones(mode_basis.grid.ndim)
			w.wcs.cdelt = mode_basis.grid.delta
			w.wcs.crval = mode_basis.grid.zero
			w.wcs.ctype = ['X', 'Y', 'Z', 'W'][:mode_basis.grid.ndim]

			hdulist[0].header.update(w.to_header())

			del tree['mode_basis']['transformation_matrix']
			tree['mode_basis']['modes'] = hdulist[0].data

		ff = asdf.fits_embed.AsdfInFits(hdulist, tree)
		ff.write_to(filename, all_array_compression='zlib', overwrite=overwrite)
	elif fmt == 'pickle':
		import pickle

		with open(filename, 'wb') as f:
			pickle.dump(mode_basis, f)
	else:
		raise NotImplementedError('The "%s" file format has not been implemented.' % fmt)
