import numpy
from io import StringIO
import re


def loadtxt2d(intext):
	try:
		return numpy.loadtxt(intext, ndmin=2)
	except:
		return numpy.loadtxt(intext)


class Analyzer(object):
	"""
		Class for accessing the output of MultiNest.

		After the run, this class provides a mean of reading the result
		in a structured way.

	"""
	
	def __init__(self, n_params, outputfiles_basename="chains/1-", verbose=True):
		self.outputfiles_basename = outputfiles_basename
		self.n_params = n_params
		"""
		[root].txt
		Compatable with getdist with 2+nPar columns. Columns have
		sample probability, -2*loglikehood, samples. Sample
		probability is the sample prior mass multiplied by its
		likelihood & normalized by the evidence.
		"""
		self.data_file = "%s.txt" % self.outputfiles_basename
		if verbose:
			print(('  analysing data from %s' % self.data_file))
		
		"""[root]post_separate.dat
		This file is only created if mmodal is set to T. Posterior
		samples for modes with local log-evidence value
		greater than nullZ, separated by 2 blank lines. Format is the
		same as [root].txt file."""
		self.post_file = "%spost_separate.dat" % self.outputfiles_basename
		
		"""[root]stats.dat
		Contains the global log-evidence, its error & local log-evidence
		with error & parameter means & standard
		deviations as well as the  best fit & MAP parameters of each of
		the mode found with local log-evidence > nullZ.
		"""
		self.stats_file = "%sstats.dat" % self.outputfiles_basename
		
		"""
		[root]post_equal_weights.dat
		Contains the equally weighted posterior samples
		"""
		self.equal_weighted_file = "%spost_equal_weights.dat" % self.outputfiles_basename
	
	def get_data(self):
		"""
			fetches self.data_file
		"""
		if not hasattr(self, 'data'):
			self.data = loadtxt2d(self.data_file)
		return self.data
	
	def get_equal_weighted_posterior(self):
		"""
			fetches self.data_file
		"""
		if not hasattr(self, 'equal_weighted_posterior'):
			self.equal_weighted_posterior = loadtxt2d(self.equal_weighted_file)
		return self.equal_weighted_posterior
	
	def _read_error_line(self, l):
		# print('_read_error_line -> line>', l)
		name, values = l.split('   ', 1)
		# print('_read_error_line -> name>', name)
		# print('_read_error_line -> values>', values)
		name = name.strip(': ').strip()
		values = values.strip(': ').strip()
		v, error = values.split(" +/- ")
		return name, float(v), float(error)
	
	def _read_error_into_dict(self, l, d):
		name, v, error = self._read_error_line(l)
		d[name.lower()] = v
		d['%s error' % name.lower()] = error
	
	def _read_table(self, txt, d=None, title=None):
		if title is None:
			title, table = txt.split("\n", 1)
		else:
			table = txt
		header, table = table.split("\n", 1)
		data = loadtxt2d(StringIO(table))
		if d is not None:
			d[title.strip().lower()] = data
		if len(data.shape) == 1:
			data = numpy.reshape(data, (1, -1))
		return data
	
	def get_stats(self):
		posterior = self.get_data()
		
		stats = []
		
		for i in range(2, posterior.shape[1]):
			b = list(zip(posterior[:, 0], posterior[:, i]))
			b.sort(key=lambda x: x[1])
			b = numpy.array(b)
			b[:, 0] = b[:, 0].cumsum()
			sig5 = 0.5 + 0.9999994 / 2.
			sig3 = 0.5 + 0.9973 / 2.
			sig2 = 0.5 + 0.95 / 2.
			sig1 = 0.5 + 0.6826 / 2.
			bi = lambda x: numpy.interp(x, b[:, 0], b[:, 1], left=b[0, 1], right=b[-1, 1])
			
			low1 = bi(1 - sig1)
			high1 = bi(sig1)
			low2 = bi(1 - sig2)
			high2 = bi(sig2)
			low3 = bi(1 - sig3)
			high3 = bi(sig3)
			low5 = bi(1 - sig5)
			high5 = bi(sig5)
			median = bi(0.5)
			q1 = bi(0.75)
			q3 = bi(0.25)
			q99 = bi(0.99)
			q01 = bi(0.01)
			q90 = bi(0.9)
			q10 = bi(0.1)
			
			stats.append({
				'median': median,
				'sigma': (high1 - low1) / 2.,
				'1sigma': [low1, high1],
				'2sigma': [low2, high2],
				'3sigma': [low3, high3],
				'5sigma': [low5, high5],
				'q75%': q1,
				'q25%': q3,
				'q99%': q99,
				'q01%': q01,
				'q90%': q90,
				'q10%': q10,
			})
		
		mode_stats = self.get_mode_stats()
		mode_stats['marginals'] = stats
		return mode_stats
	
	def get_mode_stats(self):
		"""
			information about the modes found:
			mean, sigma, maximum a posterior in each dimension
		"""
		with open(self.stats_file) as file:
			lines = file.readlines()
		# Global Evidence
		stats = {'modes': []}
		self._read_error_into_dict(lines[0], stats)
		
		if 'Nested Sampling Global Log-Evidence'.lower() in stats:
			stats['global evidence'] = stats['Nested Sampling Global Log-Evidence'.lower()]
			stats['global evidence error'] = stats['Nested Sampling Global Log-Evidence error'.lower()]
		
		if 'Nested Importance Sampling Global Log-Evidence' in lines[1]:
			# INS global evidence
			self._read_error_into_dict(lines[1], stats)
			Z = stats['Nested Importance Sampling Global Log-Evidence'.lower()]
			Zerr = stats['Nested Importance Sampling Global Log-Evidence error'.lower()]
			# use INS results in default name
			stats['global evidence'] = Z
			stats['global evidence error'] = Zerr
			
			text = ''.join(lines[3:])
			
			i = 0
			modelines = text.split("\n\n")
			mode = {
				'index': i
			}
			# no preamble; construct it
			# Strictly local evidence
			mode['Strictly Local Log-Evidence'.lower()] = Z
			mode['Strictly Local Log-Evidence error'.lower()] = Zerr
			mode['Local Log-Evidence'.lower()] = Z
			mode['Local Log-Evidence error'.lower()] = Zerr
			t = self._read_table(modelines[0], title="Parameters")
			#-----------------------------------------------------
			try:
				equ_w = self.get_equal_weighted_posterior()
				loglike = numpy.exp(equ_w.T[-1])
				#print(equ_w)
				loglikemin = loglike.min()
				loglikemax = loglike.max()
				#print(loglikemin,loglikemax)
				d_like = loglikemax-loglikemin
				#vv = loglikemax-0.954*d_like
				vv = loglikemax-0.305*d_like
				#print('vv',vv)
				good_index = numpy.where(loglike>vv)[0]
				#print('index',good_index)
				good_equ_w = equ_w[good_index][:,:-1]
				
				good_equ_w_max = good_equ_w.max(axis = 0)
				good_equ_w_min = good_equ_w.min(axis = 0)
				mean = 0.5*(good_equ_w_max+good_equ_w_min)
				sigma = 0.5*(good_equ_w_max-good_equ_w_min)
				#mean = good_equ_w.mean(axis = 0)
				#sigma = good_equ_w.std(ddof=1,axis =0)
				mode['mean'] = list(mean)
				mode['sigma'] = list(sigma)
			except:
				print('something not good!')
				mode['mean'] = t[:, 1].tolist()
				mode['sigma'] = t[:, 2].tolist()
			#-------------------------------------------------------
			mode['maximum'] = self._read_table(modelines[1])[:, 1].tolist()
			mode['maximum a posterior'] = self._read_table(modelines[2])[:, 1].tolist()
			stats['modes'].append(mode)
		else:
			text = ''.join(lines)
			# without INS:
			parts = text.split("\n\n\n")
			del parts[0]
			
			i = 0
			for p in parts:
				modelines = p.split("\n\n")
				mode = {
					'index': i
				}
				i = i + 1
				modelines1 = modelines[0].split("\n")
				# Strictly local evidence
				self._read_error_into_dict(modelines1[1], mode)
				self._read_error_into_dict(modelines1[2], mode)
				t = self._read_table(modelines[1], title="Parameters")
				mode['mean'] = t[:, 1].tolist()
				mode['sigma'] = t[:, 2].tolist()
				mode['maximum'] = self._read_table(modelines[2])[:, 1].tolist()
				mode['maximum a posterior'] = self._read_table(modelines[3])[:, 1].tolist()
				stats['modes'].append(mode)
		return stats
	
	def get_best_fit(self):
		data = self.get_data()
		i = (-0.5 * data[:, 1]).argmax()
		lastrow = data[i]
		return {'log_likelihood': float(-0.5 * lastrow[1]),
		        'parameters': list(lastrow[2:])}