#   pyRCV2: Preferential vote counting
#   Copyright © 2020  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/>.

__pragma__ = lambda x: None

from pyRCV2.method.base_stv import WIGSTVCounter
from pyRCV2.model import CandidateState, CountCard
from pyRCV2.safedict import SafeDict

class WrightSTVCounter(WIGSTVCounter):
	"""
	Wright STV implementation
	"""
	
	def exclude_candidate(self):
		"""
		Overrides BaseSTVCounter.exclude_candidate per Wright STV
		"""
		
		candidate_excluded, count_card = self.candidate_to_exclude()
		count_card.state = CandidateState.EXCLUDED
		
		# Reset the count
		# Carry over certain candidate states
		new_candidates = SafeDict()
		for candidate, count_card in self.candidates.items():
			new_count_card = CountCard()
			
			if count_card.state == CandidateState.WITHDRAWN:
				new_count_card.state = CandidateState.WITHDRAWN
			elif count_card.state == CandidateState.EXCLUDED:
				new_count_card.state = CandidateState.EXCLUDED
			
			__pragma__('opov')
			new_candidates[candidate] = new_count_card
			__pragma__('noopov')
		
		self.candidates = new_candidates
		self.exhausted = CountCard()
		self.loss_fraction = CountCard()
		self.num_elected = 0
		
		step_results = self.step_results # Carry over step results
		result = self.reset()
		self.step_results = step_results
		result.comment = 'Exclusion of ' + candidate_excluded.name
		
		return result
