"""
    Copyright 2015-2023 Stephane De Mita, Mathieu Siol

    This file is part of EggLib.

    EggLib 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.

    EggLib 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 EggLib.  If not, see <http://www.gnu.org/licenses/>.
"""
from .. import _freq

def SFS(sites, struct=None, max_missing=0.0, folded='auto', nbins=None, skip_fixed=False):
	"""
	Compute the site frequency spectrum of mutations over a list of 
	:class:`.Site` instances.
	
	
	:param sites: an iterator over :class:`.Site` instances.

	:param struct: a :class:`.Structure` instance defining the samples
        to process and/or the outgroup samples (necessary for the 
        unfolded spectrum).

    :param max_missing: maximum relative proportion of missing data in 
    the ingroup to process a site.
	
	:param folded: indicate whether the SFS should be folded or unfolded
	(possible only when outgroup is present in struct) auto -> as implied
	by the outgroup, True -> folded (even if there is an outgroup)
	False -> unfolded (error if no outgroup)
	
	:nbins: the number of bins on which to calculate the sfs. Cannot be 
	larger than number of samples/2 + 1 if folded and samples +1 if 
	unfolded. If None, the full SFS is returned as a dictionary.
	
	:skip_fixed: boolean indicating whether fixed sites should be counted
	of not.
	
	:return:A list of (bounds, counts) tuples where bounds contains the 
	top limit of each bin, which is inclusive. Except for the case where
	no bins are	provided when a dictionary containing the absolute SFS 
	is returned.
	
	"""
	
	if max_missing < 0 or max_missing > 1: raise ValueError('max_missing out of bound')
	
	#create an internal variable to determine if folded or not (once and for all)
	if folded == "auto":
		if struct is None:
			_folded = True
		else:
			if struct.num_indiv_outgroup != 0:
				_folded = False
			else:
				_folded = True
	elif folded == True:
		_folded = True
	elif folded == False:
		if struct is None:
			raise ValueError('cannot get unfolded SFS with no outgroup information')
		_folded = False

	#it = iter(sites)
	_ns = None
	freq = egglib.Freq()
	
	if nbins == None:
		counts = {}
		for site in sites:
			if not isinstance(site, egglib._site.Site):
				raise TypeError('expect a Site instance')
			
			if _ns is None:
				_ns = site.ns
				if _folded == False:
					counts = dict.fromkeys(range(_ns+1), 0) #create keys (including 0 and ns) which has to be dropped if skip_fixed is true	
				elif _folded == True:
					counts = dict.fromkeys(range(_ns//2+1), 0) #create keys (including 0) which has to be dropped if skip_fixed is true
			else:
				if sites.ns != ns:
					raise ValueError('inconsistent number of samples between sites')		
			
			freq.from_site(site, struct)
			if freq.num_alleles > 2:
				continue
			if freq.nseff()/_ns < max_missing:
				continue
			
			if _folded == False:
				#get derived allele if outgroup present
				l = (freq.freq_allele(0, cpt = egglib.Freq.outgroup), freq.freq_allele(1, cpt = egglib.Freq.outgroup))
				if 0 not in l:#polymorphic in outgroup (cannot orientate site)
					continue
				derived = 1 - l.index(max(l))
				counts[freq.freq_allele(derived, cpt = egglib.Freq.ingroup)] += 1
			elif _folded == True:
				counts[min(freq.freq_allele(0, cpt = egglib.Freq.ingroup), freq.freq_allele(1, cpt = egglib.Freq.ingroup))] += 1
				
		if skip_fixed:
			try:
				del counts[0]
				del counts[_ns]
			except KeyError:
				pass

		return counts
		
	else:
		if _folded:
			step = 0.5/nbins
		else:
			step = 1/nbins
		bounds = [i*step for i in range(1, nbins+1)]
		counts = [0 for i in range(len(bounds))]
		
		freq = egglib.Freq()
		for site in sites:
			if not isinstance(site, egglib._site.Site):
				raise TypeError('expect a Site instance')
			
			if _ns is None:
				_ns = site.ns
			else:
				if site.ns != _ns:
					raise ValueError('inconsistent number of samples between sites')
			
			freq.from_site(site, struct)
			if freq.num_alleles > 2:
				continue
			if skip_fixed and freq.num_alleles == 1:
				continue
			if freq.nseff()/_ns < max_missing:
				continue
			
			if _folded:
				minor = min(freq.freq_allele(0, cpt = egglib.Freq.ingroup), freq.freq_allele(1, cpt = egglib.Freq.ingroup))/freq.nseff()
				c = 0
				while True:
					if minor > bounds[c]:
						c += 1
					else:
						break
				counts[c] += 1
			else:
				l = (freq.freq_allele(0, cpt = egglib.Freq.outgroup), freq.freq_allele(1, cpt = egglib.Freq.outgroup))
				if 0 not in l:#polymorphic in outgroup (cannot orientate site)
					continue
				derived = 1 - l.index(max(l))
				rel = freq.freq_allele(derived, cpt = egglib.Freq.ingroup)/freq.nseff()
				c = 0
				while True:
					if rel > bounds[c]:
						c += 1
					else:
						break
				counts[c] += 1
		
		return list(zip(bounds, counts))
		
	

