#   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/>.

import pyRCV2.blt
import pyRCV2.model
import pyRCV2.numbers

from pyRCV2.method.base_stv import BaseUIGSTVCounter, BaseWIGSTVCounter
from pyRCV2.method.wright import WrightSTVCounter
from pyRCV2.ties import TiesBackwards, TiesPrompt, TiesRandom

def add_parser(subparsers):
	parser = subparsers.add_parser('stv', help='single transferable vote')
	parser.add_argument('file', help='path to BLT file')
	
	parser.add_argument('--quota', '-q', choices=['droop', 'droop_exact', 'hare', 'hare_exact'], default='droop', help='quota calculation (default: droop)')
	parser.add_argument('--quota-criterion', '-c', choices=['geq', 'gt'], default='geq', help='quota criterion (default: geq)')
	parser.add_argument('--prog-quota', action='store_true', help='progressively reducing quota')
	parser.add_argument('--numbers', '-n', choices=['fixed', 'rational', 'int', 'native'], default='fixed', help='numbers mode (default: fixed)')
	parser.add_argument('--decimals', type=int, default=5, help='decimal places if --numbers fixed (default: 5)')
	parser.add_argument('--surplus-order', '-s', choices=['size', 'order'], default='size', help='whether to distribute surpluses by size or by order of election (default: size)')
	parser.add_argument('--method', '-m', choices=['wig', 'uig', 'wright'], default='wig', help='method of surpluses and exclusions (default: wig - weighted inclusive Gregory)')
	parser.add_argument('--ties', '-t', action='append', choices=['backwards', 'prompt', 'random'], default=None, help='how to resolve ties (default: backwards then random)')

def print_step(result):
	print(result.comment)
	
	for candidate, count_card in result.candidates.items():
		state = None
		if count_card.state == pyRCV2.model.CandidateState.ELECTED or count_card.state == pyRCV2.model.CandidateState.PROVISIONALLY_ELECTED:
			state = 'ELECTED'
		elif count_card.state == pyRCV2.model.CandidateState.EXCLUDED:
			state = 'Excluded'
		elif count_card.state == pyRCV2.model.CandidateState.WITHDRAWN:
			state = 'Withdrawn'
		
		if state is None:
			print('- {}: {} ({})'.format(candidate.name, count_card.votes.pp(2), count_card.transfers.pp(2)))
		else:
			print('- {}: {} ({}) - {}'.format(candidate.name, count_card.votes.pp(2), count_card.transfers.pp(2), state))
	
	print('Exhausted: {} ({})'.format(result.exhausted.votes.pp(2), result.exhausted.transfers.pp(2)))
	print('Loss to fraction: {} ({})'.format(result.loss_fraction.votes.pp(2), result.loss_fraction.transfers.pp(2)))
	print('Total votes: {}'.format(result.total.pp(2)))
	print('Quota: {}'.format(result.quota.pp(2)))
	
	print()

def main(args):
	# Set settings
	if args.numbers == 'native':
		pyRCV2.numbers.set_numclass(pyRCV2.numbers.Native)
	elif args.numbers == 'int':
		pyRCV2.numbers.set_numclass(pyRCV2.numbers.NativeInt)
	elif args.numbers == 'rational':
		pyRCV2.numbers.set_numclass(pyRCV2.numbers.Rational)
	elif args.numbers == 'fixed':
		pyRCV2.numbers.set_numclass(pyRCV2.numbers.Fixed)
		pyRCV2.numbers.set_dps(args.decimals)
	
	with open(args.file, 'r') as f:
		election = pyRCV2.blt.readBLT(f.read())
	
	# Create counter
	if args.method == 'uig':
		counter = BaseUIGSTVCounter(election, vars(args))
	elif args.method == 'wright':
		counter = WrightSTVCounter(election, vars(args))
	else:
		counter = BaseWIGSTVCounter(election, vars(args))
	
	if args.ties is None:
		args.ties = ['backwards', 'random']
	
	counter.options['ties'] = []
	for t in args.ties:
		if t == 'backwards':
			counter.options['ties'].append(TiesBackwards())
		elif t == 'prompt':
			counter.options['ties'].append(TiesPrompt())
		elif t == 'random':
			counter.options['ties'].append(TiesRandom())
	
	# Reset
	result = counter.reset()
	
	print_step(result)
	
	# Step election
	while True:
		result = counter.step()
		
		if isinstance(result, pyRCV2.model.CountCompleted):
			break
		
		print_step(result)