# wsp-tools is TEM data analysis and simulation tools developed by WSP as a grad student in the McMorran Lab.
# Copyright (C) 2021  William S. Parker
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
from pathlib import Path
import numpy as np

__all__ = ['high_pass','low_pass','clip_data','shift_pos','outpath','ndap']

# %%
def high_pass(data, sigma = 7):
	"""Apply a high pass filter to a 2d-array.

	**Parameters**

	* **data** : _complex ndarray_ <br />

	* **sigma** : _number, optional_ <br />
	Standard deviation of the gaussian filter, measured in pixels. <br />
	Default is `sigma = 7`.

	**Returns**

	* **FFdata** : _complex ndarray_ <br />
	"""
	X = np.fft.fftfreq(data.shape[1], 1/data.shape[1])
	Y = np.fft.fftfreq(data.shape[0], 1/data.shape[0])
	x, y = np.meshgrid(X, Y)

	g = 1 - np.exp(-(x**2+y**2)/2/sigma**2)

	Fdata = np.fft.fft2(data)
	FFdata = np.fft.ifft2(g * Fdata)
	return(FFdata)

def low_pass(data, sigma = 10):
	"""Apply a low pass filter to a 2d-array.

	**Parameters**

	* **data** : _complex ndarray_ <br />

	* **sigma** : _number, optional_ <br />
	Standard deviation of the gaussian kernel (in real space), measured in pixels. <br />
	Default is `sigma = 10`.

	**Returns**

	* **FFdata** : _complex ndarray_ <br />
	"""
	X = np.linspace(-data.shape[1]/2, data.shape[1]/2, data.shape[1])
	Y = np.linspace(-data.shape[1]/2, data.shape[0]/2, data.shape[0])
	x, y = np.meshgrid(X, Y)
	g = np.exp(-(x**2 + y**2)/2/sigma**2)
	Fdata = np.fft.fft2(data)
	Fg = np.fft.fft2(np.fft.fftshift(g))
	Fg = Fg/np.max(Fg)
	FFdata = np.fft.ifft2(Fg * Fdata)
	return(FFdata)

def clip_data(data, sigma = 5):
	"""Clip data to a certain number of standard deviations from average.

	* **data** : _complex ndarray_ <br />

	* **sigma** : _number, optional_ <br />
	Number of standard deviations from average to clip to. <br />
	Default is `sigma = 5`.

	**Returns**

	* **data** : _complex ndarray_ <br />
	"""
	avg = np.mean(data)
	stdev = np.std(data)
	vmin = avg - sigma*stdev
	vmax = avg + sigma*stdev
	data[data < vmin] = vmin
	data[data > vmax] = vmax
	return(data)

def shift_pos(data):
	"""Shift data to be all greater than zero.

	**Parameters**

	* **data** : _complex ndarray_ <br />

	**Returns**

	* **data** : _complex ndarray_
	"""
	return(data - np.min(data))

def outpath(datadir, outdir, fname):
	"""A util to get the output filename.

	An example is easiest to explain:

	datadir: `/abs/path/to/data`

	fname: `/abs/path/to/data/plus/some/structure/too.dm3`

	outdir: `/where/i/want/to/write/data`

	This util will create the folder (if not exists):

	`/where/i/want/to/write/data/plus/some/structure`

	and return the filename:

	`/where/i/want/to/write/data/plus/some/structure/too`.

	**Parameters**

	* **datadir** : _string_ <br />
	The directory for the experiment's data. (abspath)

	* **outdir** : _string_ <br />
	The main directory where you want outputs to go. (abspath)

	* **fname** : _string_ <br />
	The name of the file in datadir. (abspath)

	**Returns**

	* **outname** : _string_ <br />
	The name of the file to save. (abspath)
	"""
	if not os.path.exists(outdir):
		os.makedirs(outdir)
	fname = os.path.splitext(fname)[0]
	subpath = os.path.relpath(os.path.dirname(fname), datadir)
	finoutdir = os.path.join(outdir, subpath)
	if not os.path.exists(finoutdir):
		os.makedirs(finoutdir)
	return(Path(os.path.join(finoutdir, os.path.basename(fname))))

# %%
class ndap(np.ndarray):
	"""A class that adds all the image processing methods to np.ndarray.

	The purpose of this class is just so you can write `myarray.high_pass().low_pass()` instead of `myarray = high_pass(low_pass(myarray))`.

	**Parameters**

	* **data** : _complex ndarray_ <br />
	Any type of ndarray - the methods are defined with a 2d array in mind.
	"""
	def __new__(cls, data):
		dummy = np.asarray(data, dtype=data.dtype).copy().view(cls)
		return(dummy)

	def __init__(self, data):
		self.isComplex = np.iscomplexobj(data)

	def high_pass(self, sigma = 7):
		"""Apply a high pass filter to a 2d-array.

		**Parameters**

		* **sigma** : _number, optional_ <br />
		Standard deviation of the gaussian filter, measured in pixels. <br />
		Default is `sigma = 7`.

		**Returns**

		* **FFdata** : _ndap_ <br />
		"""
		if self.isComplex:
			self[:,:] = high_pass(self, sigma)
		else:
			self[:,:] = np.real(high_pass(self, sigma))
		return(self)

	def low_pass(self, sigma = 100):
		"""Apply a low pass filter to a 2d-array.

		**Parameters**

		* **sigma** : _number, optional_ <br />
		Standard deviation of the gaussian filter, measured in pixels. <br />
		Default is `sigma = 100`.

		**Returns**

		* **FFdata** : _ndap_ <br />
		"""
		if self.isComplex:
			self[:,:] = low_pass(self, sigma)
		else:
			self[:,:] = np.real(low_pass(self, sigma))
		return(self)

	def clip_data(self, sigma = 5):
		"""Clip data to a certain number of standard deviations from average.

		* **sigma** : _number, optional_ <br />
		Number of standard deviations from average to clip to. <br />
		Default is `sigma = 5`.

		**Returns**

		* **data** : _ndap_ <br />
		"""
		self[:,:] = clip_data(self, sigma)
		return(self)

	def shift_pos(self):
		"""Shift data to be all greater than zero.

		**Returns**

		* **data** : _ndap_
		"""
		self[:,:] = shift_pos(self)
		return(self)
