#   scipy-yli: Helpful SciPy utilities and recipes
#   Copyright © 2022  Lee Yingtong Li (RunasSudo)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

from .config import config

class BayesFactor:
	"""A Bayes factor"""
	
	def __init__(self, factor, num_symbol, num_desc, denom_symbol, denom_desc):
		#: Value of the Bayes factor (*float*)
		self.factor = factor
		
		#: Symbol representing the hypothesis in the numerator (*str*)
		self.num_symbol = num_symbol
		#: Description of the hypothesis in the numerator (*str*)
		self.num_desc = num_desc
		#: Symbol representing the hypothesis in the denominator (*str*)
		self.denom_symbol = denom_symbol
		#: Description of the hypothesis in the denominator (*str*)
		self.denom_desc = denom_desc
	
	def __repr__(self):
		if config.repr_is_summary:
			return self.summary()
		return super().__repr__()
	
	def _repr_html_(self):
		return 'BF<sub>{0}{1}</sub> = {2:.2f}, {5}<br>H<sub>{0}</sub>: {3}<br>H<sub>{1}</sub>: {4}'.format(self.num_symbol, self.denom_symbol, self.factor, self.num_desc, self.denom_desc, self.interpret_lw(html=True))
	
	def summary(self):
		"""
		Return a stringified summary of the Bayes factor
		
		:rtype: str
		"""
		
		return 'BF{0}{1} = {2:.2f}, {5}\nH{0}: {3}\nH{1}: {4}'.format(self.num_symbol, self.denom_symbol, self.factor, self.num_desc, self.denom_desc, self.interpret_lw(html=False))
	
	def invert(self):
		"""
		Invert the Bayes factor
		
		:rtype: :class:`BayesFactor`
		"""
		
		return BayesFactor(1 / self.factor, self.denom_symbol, self.denom_desc, self.num_symbol, self.num_desc)
	
	def interpret_lw(self, html):
		"""
		Interpret the Bayes factor according to the Lee & Wagenmakers classification scheme
		
		:meta private:
		"""
		
		if self.factor == 1:
			return 'Evidence favours neither hypothesis to the other'
		
		if self.factor < 1:
			return self.invert().interpret_lw()
		
		if self.factor < 3:
			level = 'Anecdotal'
		elif self.factor < 10:
			level = 'Moderate'
		elif self.factor < 30:
			level = 'Strong'
		elif self.factor < 100:
			level = 'Very strong'
		else:
			level = 'Extreme'
		
		if html:
			return '{} evidence in favour of H<sub>{}</sub>'.format(level, self.num_symbol)
		else:
			return '{} evidence in favour of H{}'.format(level, self.num_symbol)

def bayesfactor_afbf(params, cov, n, hypothesis):
	"""
	Compute an adjusted fractional Bayes factor for the hypothesis
	
	Uses R *BFpack* library. See R documentation for *BFpack.BF*.
	
	:return: Bayes factor for the hypothesis versus its complement
	:rtype: :class:`yli.bayes_factors.BayesFactor`
	"""
	
	import rpy2.robjects as ro
	import rpy2.robjects.packages
	import rpy2.robjects.pandas2ri
	
	# Import BFpack
	ro.packages.importr('BFpack')
	
	with ro.conversion.localconverter(ro.default_converter + ro.pandas2ri.converter):
		with ro.local_context() as lc:
			lc['params'] = params
			lc['cov'] = cov
			lc['n'] = n
			lc['hypothesis'] = hypothesis
			
			ro.r('bf_fit <- BF(params, Sigma=as.matrix(cov), n=n, hypothesis=hypothesis)')
			bf_matrix = ro.r('bf_fit$BFmatrix_confirmatory')
			
			return BayesFactor(bf_matrix[0][1], '1', hypothesis, 'C', 'Complement')