import os
import os.path
import urllib2
import urlparse
import hashlib
import gzip

import numpy as np
import scipy

from pysit import *
from pysit.util.config import get_gallery_data_path
from pysit.util.net import download_file
from pysit.util.io import read_model
from pysit.util.image_processing import blur_image

__all__ = ['GalleryModel', 'GeneratedGalleryModel', 'PrecomputedGalleryModel']

class GalleryModel(object):
	""" Base class for gallery models in PySIT.
	
	Attribtues
	----------
	model_name : string, class attribute
		Long form name of the model.
	valid_dimensions : iterable, class attribute
		List of the dimensions for which the model is valid.
	supported_physics : list, class attribute
		List of the physics (e.g., acoustic, or elastic) supported by the model.
	physics_to_parameters_map : dict, class attribute
		Dict which maps physics to the list of parameters needed for that physics.
	dimension : int
		Dimension of the actual instance of the gallery model.
	true_model : ndarray
		The true model for a gallery instance.
	initial_model : ndarray
		The initial model for a gallery instance.
	mesh : pysit.mesh
		The computational mesh on which the model is defined.
	domain : pysit.domain
		The physical domain on which the model is defined.
		
	"""
	

	model_name = ""
		
	valid_dimensions = tuple()
		
	# List of physics supported for this gallery problem
	supported_physics = list()
	
	physics_to_parameters_map = {'acoustic' : ['vp'],
		                         'variable-density-acoustic' : ['vp', 'density'],
	                             'elastic' : ['vp', 'vs', 'density']}

	@property #read only
	def dimension(self):
		raise NotImplementedError('dimension must be set in subclass')
		
	def __init__(self, *args, **kwargs):
		self._true_model = None
		self._initial_model = None
		self._mesh = None
		self._domain = None
		
	def get_setup(self):
		""" Return the usual 4-tuple of (true_model, initial_model, mesh, and 
		    domain) needed for an experiment."""
		
		true_model = self.true_model
		
		initial_model = self.initial_model
		
		mesh = self.mesh
		
		domain = self.domain
		
		return true_model, initial_model, mesh, domain
	
	@property
	def true_model(self):
		return self._true_model
	
	@property
	def initial_model(self):
		return self._initial_model
	
	@property
	def mesh(self):
		return self._mesh
		
	@property
	def domain(self):
		return self._domain
		
class GeneratedGalleryModel(GalleryModel):
	""" Empty base class for generated gallery models."""
	pass
		
		
		
class PrecomputedGalleryModel(GalleryModel):
	""" Base class for pre-computed or community gallery models in PySIT.
	
	Precomputed models are licensed (or open) community models.  Given that
	distribution licensing is tricky, this class is designed to download any
	model that is not already downloaded to the location specified in the pysit
	configuration file.  In this manner, the user is responsible for ensuring
	that they have the proper approval to use the community model.
	
	Any rescalings and crops are cached in the user's specified cache folder to
	speed up computation.
	
	Attribtues
	----------
	fs_full_name : string, class attribute
		Longer name for use in file system calls.
	fs_short_name : string, class attribute
		Short name for use in file system calls.
	supported_physical_parameters : iterable, class attribute
		List of the physical parameters supported by the gallery model.
	supported_masks : iterable, class attribute
		List of the masks supported by the gallery model.
	base_physical_origin : iterable, class attribute
		Origin of the full sized model coordinate system.
	base_physical_size : iterable, class attribute
		Size (in physical units) of the full sized model coordinate system.
	base_physical_dimensions_units : list, class attribute
		String with the spatial units for the coordinate system.
	base_pixels : iterable, class attribute
		Number of pixels in the full sized model.
	base_pixel_scale : iterable, class attribute
		Size of the pixels in the full sized model.
	base_pixel_units : list, class attribute
		Units for the the pixels in the full sized model.
	water_properties : tuple, class attribute
		Specifies how water layers are determined in the model:
			* (None, ) indicates no water
			* ('depth', <depth: float>) specifies that a preset depth from the base_physical_origin is to be used
			* ('mask', ) indicates that the gallery model's water mask should be used
	patches : dict, class attribute
		Mapping of pre-specified useful patches of the gallery model to their origin and sizes.
	
	Methods
	-------
	rebuild_models : 
		Rebuilds true and initial models if changes are made to the underlying specification.
	
	"""

	
	# A sanitized name for filesystem work
	fs_full_name = ""
	fs_short_name = ""
	
	# Available data
	supported_physical_parameters = list() #['density', 'vp']
	supported_masks = list() # ['salt', 'water']

	# The local file names used to save the downloaded models
	_local_parameter_filenames = {} #{'vp' : None, 'vs' : None, 'density' : None}
	_local_mask_filenames = {} #{'water-mask' : None, 'salt-mask' : None}
	
	# Sometimes the sources are not in the units we expect.  This scale factor
	# allows the numbers to be converted into common units.
	_parameter_scale_factor = {} #{ 'vp' : 1.0, 'density' : 1.0}
		                        		                   
	# Dictionary mapping physical parameters to a list of URLs where they
	# can be obtained   	
	_remote_file_sources = {}
	
	# Sometimes the model data is stored transposed from how pysit expects it.
	_model_transposed = False
	
	# Model specification
	base_physical_origin =  None
	base_physical_size =None
	base_physical_dimensions_units = tuple()
	
	base_pixels = None
	base_pixel_scale = None
	base_pixel_units = tuple()

	# Water properties specify the way the water masking is handled
	water_properties = (None, ) 
	
	# Default configuration of the four possible initial model schemes.
	_initial_configs = {'smooth_width': {'sigma':3000.0},
		                'smooth_low_pass': {'freq':1/3000.},
		                'constant': {'velocity': 3000.0},
		                'gradient': {'min':1500.0, 'max':3000}}
	
	# Maps names of scales to physical scales for the model
	_scale_map = {}		                	

	# Common array patches		          	
	patches = {} 
	
	# Short cuts for the left and right positions of the edges of the domain.
	@property
	def _coordinate_lbounds(self): return self.physical_origin
	@property
	def _coordinate_rbounds(self): return self.physical_origin + self.physical_size

	license_string = None


	def __init__(self, physics='acoustic', 
	                   origin=None,
	                   size=None,
	                   pixel_scale='mini', 
	                   pixels=None,
	                   initial_model_style='smooth_low_pass',
	                   initial_config={},
	                   fix_water_layer=True,	                   	
	                   **kwargs):
	                   	
		""" Constructor for community gallery models.
		
		Parameters
		----------
		physics : {'acoustic', 'variable-density-acoustic', 'elastic'}
			Wave physics to instantiate.
		origin : array-like, optional
			Origin, in physical coordinates, of the model patch.  Defaults to base origin.
		size : array-like, optional
			Size, in physical units, of the model patch.  Defaults to base size.
		pixel_scale : string or array-like
			Either key to pixel scale map, or array-like container of grid spacings
		pixels : array-like
			Number of pixels. This overrides pixel_scale.
		initial_model_style: {'smooth_low_pass', 'smooth_width', 'constant', 'gradient'}
			Style of the initial model, see note for more details.
		initial_config : dict
			Configuration of the initial model style
		fix_water_layer : boolean
			If true, water layer in initial model is set to match water layer in true model.
		kwargs : dict
			Additional keyword arguments for specifying boundary conditions and more.	
		
		
		Notes
		-----
		The four initial model styles are:
		1) 'smooth_low_pass' -- Smooths the true model with a Gaussian low pass filter with specified cutoff frequency.
			Options: 'freq' : spatial frequency for the cutoff
		2) 'smooth_width' -- Smooths the true model with a Gaussian low pass filter with specified standard deviation.
			Options: 'sigma' : width of the smoothing filter
		3) 'constant' -- Initial value is a specified constant
			Options: 'velocity' : the velocity to use
		4) 'gradient' -- Sets initial velocity to a linear gradient
			Options: 'min' : Minimum value for the gradient.
			         'max' : Maximum value for the gradient.
		
		"""
		
		# Validate physics and collect the names of the associated parameters.
		if physics not in self.supported_physics:
			raise ValueError('Physics \'{0}\' not supported by model {1}'.format(physics, self.__class__.__name__))
		
		# Currently only support accoustic physics.
		if physics != 'acoustic':
			raise NotImplementedError('Currently only support accoustic physics.')
		
		self.physics = physics
		
		# If origin is not set, use the model base.  Also verify that it is within valid bounds.
		if origin is not None:
			if np.all(origin >= self.base_physical_origin) and np.all(origin < (self.base_physical_size - self.base_physical_origin)):
				self.physical_origin = np.asarray(origin).reshape(2,)
			else:
				raise ValueError('Specified origin is outside of the model bounds.')
		else:
			self.physical_origin = self.base_physical_origin.copy() 
		
		# If size is not set, use the model base.  Also verify that it is within valid bounds.
		if size is not None:
			size = np.asarray(size).reshape(2,)
			if np.all(size > 0.0) and np.all((size + self.physical_origin) <= self.base_physical_size):
				self.physical_size = size
			else:
				raise ValueError('Specified size is outside of the model bounds.')
		else:
			self.physical_size = self.base_physical_size.copy() 
			
		# Setting the number of pixels trumps the pixel_scale 
		if pixels is not None:
			self.pixels = np.asarray(pixels).reshape(2,)
			self.pixel_scale = self.physical_size / pixels
			
		elif pixel_scale is not None:
			if isinstance(pixel_scale, str):
				pixel_scale = self._scale_map[pixel_scale]
				
			self.pixel_scale = np.asarray(pixel_scale).reshape(2,)
			
			length = self._coordinate_rbounds - self._coordinate_lbounds
			self.pixels = np.ceil(length/self.pixel_scale).astype(np.int)
			
			# correct physical size to match the number of pixels.  this can go over the true max size by up to 1 delta
			self.physical_size = self.pixel_scale*self.pixels
			
			# Add 1 because the last point _is_ included (but not yet)
#			self.pixels += 1
		
		else: #both are None
			self.pixels = self.base_pixels.copy()
			self.pixel_scale = self.base_pixel_scale.copy()
		
		# Create pysit domains and meshes
		
		# extract or default the boundary conditions
		x_lbc = kwargs['x_lbc'] if ('x_lbc' in kwargs) else PML(0.1*self.pixels[0]*self.pixel_scale[0], 100.0)
		x_rbc = kwargs['x_rbc'] if ('x_rbc' in kwargs) else PML(0.1*self.pixels[0]*self.pixel_scale[0], 100.0)
		z_lbc = kwargs['z_lbc'] if ('z_lbc' in kwargs) else PML(0.1*self.pixels[1]*self.pixel_scale[1], 100.0)
		z_rbc = kwargs['z_rbc'] if ('z_rbc' in kwargs) else PML(0.1*self.pixels[1]*self.pixel_scale[1], 100.0)
		
		x_config = (self._coordinate_lbounds[0], self._coordinate_rbounds[0], x_lbc, x_rbc)
		z_config = (self._coordinate_lbounds[1], self._coordinate_rbounds[1], z_lbc, z_rbc)
		d = RectangularDomain(x_config, z_config)
		
		mesh = CartesianMesh(d, *self.pixels)
		
		# store initial model setup values
		self.initial_model_style = initial_model_style
		self.initial_config = initial_config
		self.fix_water_layer = fix_water_layer
		
		# Store the mesh and domain, and build the model parameters
		self._mesh = mesh
		self._domain = mesh.domain
		self.rebuild_models()
		
	def rebuild_models(self):
		""" Rebuild the true and initial models based on the current configuration."""
		
		# Get the list of parameters required for the current physics
		parameter_names = self.physics_to_parameters_map[self.physics]
		
		# Get a list of properly scaled and cropped parameters arrays for each of those parameters
		parameter_arrays = [self._load_scaled_parameter_array(p) for p in parameter_names]
		
		# Velocity only, at the moment
		# Eventually all of this will need to be done for each model parameter
		vp = parameter_arrays[0]
		
		# Extract the depth coordinates of the grid points
		grid = self._mesh.mesh_coords()
		ZZ = grid[-1].reshape(vp.shape)
		
		# Copy the default values for the initial model style
		config = dict(self._initial_configs[self.initial_model_style])
		# Override them with user provided values
		config.update(self.initial_config)
		
		# Construct initial velocity profile:
		if self.initial_model_style == 'constant':
			vp0 = np.ones(vp.shape)*config['velocity']
						
		elif self.initial_model_style == 'smooth_width': 
			vp0 = blur_image(vp, 
			                 sigma=config['sigma'], 
			                 mesh_deltas=self._mesh.deltas)
						
		elif self.initial_model_style == 'smooth_low_pass': 
			vp0 = blur_image(vp, 
			                 freq=config['freq'], 
			                 mesh_deltas=self._mesh.deltas)
			
		elif self.initial_model_style == 'gradient':		
			# construct the N-D depth trace shape, which is 1 in every dimension
			# except the depth direction
			sh = self._mesh.shape(as_grid=True)
			_shape_tuple = tuple([1]*(len(sh)-1) + [sh[-1]])
			
			# Construct the shape of the padding, which is 0 padding in the depth
			# direction and n pixels in each other direction
			_pad_tuple = [(0,n-1) for n in sh]
			_pad_tuple[-1] = (0,0)
			_pad_tuple = tuple(_pad_tuple)
			
			# If the water is to be fixed, the gradient is started at the 
			# shallowest water-rock interaction. Otherwise, start at the top.
			if self.fix_water_layer and (self.water_properties[0] is not None):
				
				grid_sparse = self._mesh.mesh_coords(sparse=True)
				ZZZ = grid_sparse[-1].reshape(_shape_tuple)
				
				min_max_depth = self._find_min_max_water_depth()
				
				loc = np.where(ZZZ > min_max_depth)
				zprof = np.zeros(_shape_tuple)
				zprof[loc] = np.linspace(config['min'],config['max'],loc[0].size)
			else:
				zprof = np.linspace(config['min'],config['max'],self._mesh.z.n)
				zprof.shape = _shape_tuple
				
			# Pad out the gradient.
			vp0 = np.pad(zprof, _pad_tuple, 'edge')
			
		if self.fix_water_layer and (self.water_properties[0] is not None):
		
			# If the water layer is specified by depth, reset everything shallower
			if self.water_properties[0] == 'depth':
				loc = np.where(ZZ < self._find_min_max_water_depth())
				vp0[loc] = vp[loc]
			
			# If the water mask is set, reset the masked values.
			elif self.water_properties[0] == 'mask':
				mask = self._load_scaled_parameter_array('water')
				self._water_mask = mask
				loc = np.where(mask != 0.0)
				vp0[loc] = vp[loc]
			
		# Construct final padded velocity profiles
		C = vp.reshape(self._mesh.shape())
		C0 = vp0.reshape(self._mesh.shape())
		
		self._true_model = C
		self._initial_model = C0	
	
	def _find_min_max_water_depth(self):
		""" Returns the shallowest water-rock interation depth. """
		
		if self.water_properties[0] == 'depth':
			return self.water_properties[1]
		
		elif self.water_properties[0] == 'mask':
			mask = self._load_scaled_parameter_array('water')
		
			grid = self._mesh.mesh_coords()
			ZZ = grid[-1].reshape(mask.shape)
			
			return np.min((ZZ*mask).max(axis=-1))
		else:
			return 0.0
		
	@classmethod
	def _get_gallery_model_path(cls):
		""" Returns the full path of the gallery data location."""
		gdp = get_gallery_data_path()
		path = os.path.join(gdp, cls.fs_full_name)
		
		if not os.path.exists(path):
			os.mkdir(path)
		return path

	@classmethod
	def _build_local_parameter_path(cls, param):
		""" Builds the full path to the local parameter or mask file."""
		
		file_path = cls._get_gallery_model_path()
		
		if param in cls.supported_physical_parameters:
			parameter_fname = cls._local_parameter_filenames[param]
		elif param in cls.supported_masks:
			parameter_fname = cls._local_mask_filenames[param]
			
		return os.path.join(file_path, parameter_fname)			

	@classmethod
	def _build_parameter_patch_filename(cls, param, hash_args, ext=''):
		""" Builds filename for rescaled model parameters of the form <fs_short_name>_<md5_str>.<ext>."""
		
		hash_str = hashlib.md5(''.join([a.tostring() for a in hash_args])).hexdigest()
		
		return "{0}_{1}_{2}{3}".format(cls.fs_short_name, param, hash_str, ext)
				
	@classmethod
	def _build_parameter_patch_path(cls, param, hash_args, ext=''):
		""" Builds the full path to the local patch file."""
		
		file_path = cls._get_gallery_model_path()
		
		parameter_fname = cls._build_parameter_patch_filename(param, hash_args, ext)
			
		return os.path.join(file_path, parameter_fname)		
				
	@classmethod
	def _write_npy(cls, arr, fname, gz=True):
		""" Writes the array as a numpy file (optionally gzipped)."""
		
		if fname.endswith('.gz'):
			f = fname[:-3]
			np.save(f, arr)
			with open(f, 'rb') as raw:
				with gzip.open(fname, 'wb') as gzz:
					gzz.write(raw.read())
			os.remove(f)
		elif gz:
			np.save(fname, arr)
			with open(fname, 'rb') as raw:
				with gzip.open(fname.join('.gz'), 'wb') as gzz:
					gzz.write(raw.read())
			os.remove(f)
		else:
			np.save(fname, arr)		
				
	@classmethod
	def _read_npy(cls, fname):
		""" Reads a potentially gzipped numpy array file.  Gets around any 
		inconsistencies in the numpy file writing specification."""
		
		if fname.endswith('.gz'):
			f = fname[:-3]
			with gzip.open(fname, 'rb') as gzz:
				arr = np.load(gzz)
		else:
			arr = np.load(fname)	
			
		return arr	
		
	@classmethod
	def _download_and_prepare(cls, param):
		""" Downloads and prepares parameter file for use.  Prepare means it reads the segy file and saves it as a gzipped numpy file."""
		
		# verify that the requested parameter is supported by this gallery model
		if param not in cls.supported_physical_parameters + cls.supported_masks:
			raise ValueError('{0} is invalid physical parameter/mask for gallery model {1}.'.format(param, cls.__name__))
		
		parameter_path = cls._build_local_parameter_path(param)
		
		# Check if file exists locally
		if not os.path.isfile(parameter_path):
		
			# If it does not, try downloading it from sources
			downloaded = False
			
			if cls.license_string is not None:
				
				print cls.license_string
				agree = raw_input("Continue downloading {0} model '{1}' (yes/[no])?  ".format(cls.model_name, param) ) or 'no'
				if agree.strip().lower() != 'yes':
					raise Exception('Cannot proceed without license agreement.')
			
			for url in cls._remote_file_sources[param]:
				
				try:
					download_file(url, parameter_path)
					downloaded = True
				except urllib2.HTTPError:
					continue
		
			# If no source is successful, raise an exception
			if not downloaded:
				raise Exception('Unable to download parameter.')
		
		# Now check and see if the 'full sized' patch version exists.  If it does not, build that.
		hash_args = (cls.base_physical_origin, cls.base_physical_size, cls.base_pixel_scale, cls.base_pixels)		
		scaled_parameter_path = cls._build_parameter_patch_path(param, hash_args, '.npy.gz')
		
		# Check if file exists locally
		if not os.path.isfile(scaled_parameter_path):
			# create it by loading the segy file and saving
			
			# If the parameter file is gzipped, unzip it
			unpacked = False
			if parameter_path[-2:] == 'gz':
				unpacked_parameter_path = parameter_path[:-3]
					
				with gzip.open(parameter_path, 'rb') as infile:
					with open(unpacked_parameter_path, 'wb') as outfile:
						outfile.write(infile.read())
				
				parameter_path = unpacked_parameter_path
				unpacked = True
			
			arr = read_model(parameter_path)
			
			# Masks are not scaled
			if param not in cls.supported_masks:
				arr *= cls._parameter_scale_factor[param]
			
			if cls._model_transposed:
				arr = arr.T
				
			# Write it
			cls._write_npy(arr, scaled_parameter_path)
			
			if unpacked:
				# delete the unzipped version
				os.remove(parameter_path)	

	def _resample_parameter(self, param):
		""" Resamples the specified parameter to the sizes, from the original."""
		
		# Load the original parameter array
		base_hash_args = (self.base_physical_origin, self.base_physical_size, self.base_pixel_scale, self.base_pixels)		
		base_scaled_parameter_path = self._build_parameter_patch_path(param, base_hash_args, '.npy.gz')
		big_arr = self._read_npy(base_scaled_parameter_path)
		
		# Resample the array, in each dimension, with nearest neighbor interpolation.
		y = big_arr
		for i in xrange(big_arr.ndim):
			
			x = self.base_physical_origin[i] + np.arange(self.base_pixels[i])*self.base_pixel_scale[i]
			I = scipy.interpolate.interp1d(x, y, copy=False, axis=i, kind='nearest')
			
			new_x = self.physical_origin[i] + np.arange(self.pixels[i])*self.pixel_scale[i]
			
			# Bounds that go slightly past the true thing should extend from the right side
			new_x[new_x > self.base_physical_size[i]] = self.base_physical_size[i]
			
			y = I(new_x)
			
		return y
	
				
	def _load_scaled_parameter_array(self, param):
		""" Builds or loads a numpy array for the scaled version of the model file."""
		
		# Download and prep the file, if it is needed
		self._download_and_prepare(param)
		
		# Build scaled filename and its path
		hash_args = (self.physical_origin, self.physical_size, self.pixel_scale, self.pixels)		
		scaled_parameter_path = self._build_parameter_patch_path(param, hash_args, '.npy.gz')
		
		# if file exists, load it
		if os.path.isfile(scaled_parameter_path):
			arr = self._read_npy(scaled_parameter_path)
		else:
			# otherwise we need to write that particular patch, as loaded from the default array.
			arr = self._resample_parameter(param)
			self._write_npy(arr, scaled_parameter_path)
		
		return arr

				
			
	@classmethod
	def _save_model(cls, format='mat'): #'npy', 'binary', 'pickle', 'segy'?
		pass

#	def load_data(self,fnames):
#		
#		models = dict()
#		
#		for model_name, fname in fnames.items():
#			m = spio.loadmat(fname)
#			
#			# correct for some savemat stupidity
#			m['pixel_scale'].shape = -1,
#			m['pixel_unit'].shape =  -1,
#			m['domain_size'].shape =  -1,
#			m['domain_units'].shape =  -1,
#			m['scale_factor'].shape =  -1,
#			m['sz'].shape =  -1,
#			
#			models[model_name] = m
#		
#		return models