# Code for reading in and analyzing outputs of imfit

import glob
import numpy as np

from . import imfit_funcs as imfuncs

# If pandas is available, use its (much faster) file-reading code
usingPandas = False
try:
	import pandas as pd
	usingPandas = True
	def GetDataColumns( filename, usecols=None ):
		df = pd.read_csv(filename, delim_whitespace=True, comment='#', dtype=np.float64,
						usecols=usecols, header=None)
		return df.values
except:
	def GetDataColumns( filename, usecols=None ):
		return np.loadtxt(filename, usecols=usecols)


# dictionary mapping imfit function short names (as found in the config/parameter file) to
# corresponding 1-D Python functions in imfit_funcs.py, along with some useful information:
#    "function" = corresponding imfit_funcs.py function, if one exists
#    "nSkip" = the number of 2D-related  parameters to skip (e.g., PA, ellipticity), 
#    "ell" = index for ellipticity parameter, if it exists, 
#    "a" = index or indices for semi-major-axis parameters (r_e, h, sigma, etc.)
# THIS IS NOT MEANT TO BE A COMPLETE LIST
imfitFunctionMap = {"Exponential": {"function": imfuncs.Exponential, "nSkip": 2, "ell": 1, "a": [3]}, 
				"Exponential_GenEllipse": {"function": imfuncs.Exponential, "nSkip": 3, "ell": 1, "a": [4]},
				"Sersic":  {"function": imfuncs.Sersic, "nSkip": 2, "ell": 1, "a": [4]},
				"Sersic_GenEllipse":  {"function": imfuncs.Sersic, "nSkip": 3, "ell": 1, "a": [5]}, 
				"Gaussian":  {"function": imfuncs.Gauss, "nSkip": 2, "ell": 1, "a": [3]},
				"GaussianRing":  {"function": imfuncs.GaussRing, "nSkip": 2, "ell": 1, "a": [3,4]},
				"GaussianRing2Side":  {"function": imfuncs.GaussRing2Side, "nSkip": 2, "ell": 1, "a": [3,4,5]},
				"Moffat": {"function": imfuncs.Moffat, "nSkip": 2, "ell": 1, "a": [3]},
				"BrokenExponential":  {"function": imfuncs.BrokenExp, "nSkip": 2, "ell": 1, "a": [3,4,5]}}



def ChopComments( theLine ):
	return theLine.split("#")[0]


def GetFunctionImageNames( baseName, funcNameList ):
	"""Generate a list of FITS filenames as would be created by makeimage in "--output-functions"
	mode.

	Parameters
	----------
	baseName : str
		root name of output files

	funcNameList : list of str
		list containing function names (e.g., ["Exponential", "Sersic", "Sersic"]

	Returns
	-------
	imageNameList : list of str
		list of output filenames
	"""
	
	nImages = len(funcNameList)
	imageNameList = [ "%s%d_%s.fits" % (baseName, i + 1, funcNameList[i]) for i in range(nImages) ]
	return imageNameList


def ReadImfitConfigFile( fileName, minorAxis=False, pix=1.0, getNames=False, X0=0.0 ):
	"""Function to read and parse an imfit-generated parameter file (or input config file)
	and return a tuple consisting of: (list of 1-D imfit_funcs functions, list of lists of parameters).
	
	pix = scale in arcsec/pixel, if desired for plotting vs radii in arcsec.
	
	We assume that all functions have a center at x = 0; this can be changed by setting
	X0.
	
	Returns tuple of (functionList, trimmedParameterList)
	If getNames == True:
		Returns tuple of (functionNameList, functionList, trimmedParameterList)

	Parameters
	----------
	fileName : str
		Imfit configuration or best-fit parameter file

	minorAxis : bool, optional

	pix : float, optional
		scale in arcsec/pixel, if desired for plotting radii in arcsec

	getNames : bool, optional
		if True, output is tuple of (functionNameList, functionList, trimmedParameterList)

	X0 : float, optional
		value for function centers

	Returns
	-------
	(functionList, trimmedParameterList) : (list of image-functions, list of float)
	"""
	
	dlines = [ line for line in open(fileName) if len(line.strip()) > 0 and line[0] != "#" ]

	funcNameList = []
	paramMetaList = []
	currentParamList = []
	nLines = len(dlines)
	for line in dlines:
		trimmedLine = ChopComments(line)
		if trimmedLine.find("X0") == 0:
			continue
		if trimmedLine.find("Y0") == 0:
			continue
		if trimmedLine.find("FUNCTION") == 0:
			# if this isn't the first function, store the previous set of parameters
			if len(currentParamList) > 0:
				paramMetaList.append(currentParamList)
			# make a new parameter list for the new function
			currentParamList = [X0]
			pp = trimmedLine.split()
			fname = pp[1].strip()
			funcNameList.append(fname)
			continue
		else:
			pp = trimmedLine.split()
			newValue = float(pp[1])
			currentParamList.append(newValue)

	# ensure that final set of parameters get stored:
	paramMetaList.append(currentParamList)
	
	# process function list to remove unneeded parameters (and convert size measures
	# from major-axis to minor-axis, if requested)
	funcList = [ imfitFunctionMap[fname]["function"] for fname in funcNameList ]
	trimmedParamList = []
	nFuncs = len(funcList)
	for i in range(nFuncs):
		fname = funcNameList[i]
		nSkipParams = imfitFunctionMap[fname]["nSkip"]
		fullParams = paramMetaList[i]
		# calculate scaling factor for minor-axis values, if needed
		if minorAxis is True:
			ellIndex = imfitFunctionMap[fname]["ell"]
			ell = fullParams[ellIndex+1]
			q = 1.0 - ell
		else:
			q = 1.0
		smaIndices = imfitFunctionMap[fname]["a"]
		# convert length values to arcsec and/or minor-axis, if needed,
		for smaIndex in smaIndices:
			# +1 to account for X0 value at beginning of parameter list
			fullParams[smaIndex+1] = pix*q*fullParams[smaIndex+1]
		# construct the final 1-D parameter set for this function: X0 value, followed
		# by post-2D-shape parameters
		trimmedParams = [fullParams[0]]
		trimmedParams.extend(fullParams[nSkipParams+1:])
		trimmedParamList.append(trimmedParams)
		
	
	if getNames is True:
		return (funcNameList, funcList, trimmedParamList)
	else:
		return (funcList, trimmedParamList)




# Code for reading output of bootstrap resampling and MCMC chains

def GetBootstrapOutput( filename ):
	"""Reads imfit's bootstrap-resampling output when saved using the
	--save-bootstrap command-line option.

	Parameters
	----------
	filename : str
		name of file with bootstrap-resampling output
	
	Returns
	-------
	(column_names, data_array) : tuple of (list, np.ndarray)
		column_names = list of column names (strings)
		data_array = numpy array of parameter values 
			with shape = (n_iterations, n_parameters)
	"""
	
	# get first 100 lines (or all lines in file, whichever is shorter)
	firstLines = []
	with open(filename) as theFile:
		try:
			for x in range(100):
				firstLines.append(next(theFile))
		except StopIteration:
			pass

	# find header line with column names and extract column names
	for i in range(len(firstLines)):
		if firstLines[i].find("# Bootstrap resampling output") >= 0:
			columnNamesIndex = i + 1
			break
	columnNames = firstLines[columnNamesIndex][1:].split()
	for i in range(len(columnNames)):
		if columnNames[i] == "likelihood":
			nParamColumns = i
			break
	
	# get the data
	d = GetDataColumns(filename)
	
	return (columnNames, d)



def GetSingleChain( filename, getAllColumns=False ):
	"""Reads a single MCMC chain output file and returns a tuple of column names
	and a numpy array with the data.
	
	Parameters
	----------
	filename : str
		name of file with MCMC output chain
	
	getAllColumns: bool, optional
		if False [default], only model parameter-value columns are retrieved;
		if True, all output columns (including MCMC diagnostics) are retrieved
	
	Returns
	-------
	(column_names, data_array) : tuple of (list, np.ndarray)
		column_names = list of column names (strings)
		data_array = numpy array of parameter values 
			with shape = (n_iterations, n_parameters)
	"""
	
	# get first 100 lines
	# FIXME: file *could* be shorter than 100 lines; really complicated
	# model could have > 100 lines of header...
	with open(filename) as theFile:
		firstLines = [next(theFile) for x in range(100)]

	# find header line with column names and extract column names
	for i in range(len(firstLines)):
		if firstLines[i].find("# Column Headers") >= 0:
			columnNamesIndex = i + 1
			break
	columnNames = firstLines[columnNamesIndex][1:].split()
	for i in range(len(columnNames)):
		if columnNames[i] == "likelihood":
			nParamColumns = i
			break
	
	# get data for all columns, or just the model parameters?
	whichCols = None
	if not getAllColumns:
		whichCols = list(range(nParamColumns))
		outputColumnNames = columnNames[:nParamColumns]
	else:
		whichCols = None
		outputColumnNames = columnNames
	
	# get the data
	d = GetDataColumns(filename, usecols=whichCols)
	
	return (outputColumnNames, d)


def MergeChains( fname_root, maxChains=None, getAllColumns=False, start=10000, last=None,
				secondHalf=False  ):
	"""
	Reads and concatenates all MCMC output chains with filenames = fname_root.*.txt,
	using data from t=start onwards. By default, all generations from each chain
	are extracted; this can be modified with the start, last, or secondHalf keywords.
	

	Parameters
	----------
	fname_root : str
		root name of output chain files (e.g., "mcmc_out")
	
	maxChains : int or None, optional
		maximum number of chain files to read [default = None = read all files]
	
	getAllColumns : bool, optional
		if False [default], only model parameter-value columns are retrieved;
		if True, all output columns (including MCMC diagnostics) are retrieved

	start : int, optional
		extract samples from each chain beginning with time = start
		ignored if "secondHalf" is True or if "last" is not None

	last : int or None, optional
		extract last N samples from each chain
		ignored if "secondHalf" is True
	
	secondHalf : bool, optional
		if True, only the second half of each chain is extracted
		if False [default], 
		
	Returns
	-------
	(column_names, data_array) : tuple of (list, np.ndarray)
		column_names = list of column names (strings)
		data_array = numpy array of parameter values 
			with shape = (n_samples, n_parameters)
	"""
	
	# construct list of filenames
	if maxChains is None:
		globPattern = "{0}.*.txt".format(fname_root)
		filenames = glob.glob(globPattern)
	else:
		filenames = ["{0}.{1}.txt".format(fname_root, n) for n in range(maxChains)]
	nFiles = len(filenames)
	
	if (nFiles < 1):
		print("ERROR: No MCMC output files found using pattern \"{0}\"".format(globPattern))
		return None
	
	# get the first chain so we can tell how long the chains are
	(colNames, dd) = GetSingleChain(filenames[0], getAllColumns=getAllColumns)
	nGenerations = dd.shape[0]

	# figure out what part of full chain to extract
	if secondHalf is True:
		startTime = int(np.floor(nGenerations / 2))
	elif last is not None:
		startTime = -last
	else:
		startTime = start
	
	# get first chain and column names; figure out if we get all columns or just
	# model parameters
	if (startTime >= nGenerations):
		txt = "WARNING: # generations in MCMC chain file {0} ({1:d}) is <= ".format(filenames[0],
				nGenerations)
		txt += "requested start time ({0:d})!\n".format(startTime)
		print(txt)
		return None
	dd_final = dd[startTime:,:]
	if getAllColumns is False:
		nParamColumns = len(colNames)
		whichCols = list(range(nParamColumns))
	else:
		whichCols = None

	# get and append rest of chains if more than 1 chain-file was requested
	if nFiles > 1:
		for i in range(1, nFiles):
			dd_next = GetDataColumns(filenames[i], usecols=whichCols)
			dd_final = np.concatenate((dd_final, dd_next[startTime:,:]))

	return (colNames, dd_final)

