#!/usr/bin/python

# Public domain software - no copyright
# Written by Jack Whitham, December 2008.
# http://www.jwhitham.org.uk/

import random, sys
try:
    import tty, termios
except:
    tty = termios = None


DESCRIPTION = {
    (0, 1) : "the defenders lose one army",
    (1, 0) : "the attackers lose one army",
    (0, 2) : "the defenders lose two armies",
    (2, 0) : "the attackers lose two armies",
    (1, 1) : "both sides lose an army",
    }


class Engine:
    def __init__(self, num_attacker_armies, num_defender_armies, verbose):
        self.num_attacker_armies = num_attacker_armies
        self.num_defender_armies = num_defender_armies
        self.old_style_rules = True
        self.verbose = verbose

    def Produce(self, num_dice):
        assert 1 <= num_dice <= 3
        dice = [ random.randint(1, 6) for i in xrange(num_dice) ]
        # Sort into descending order
        dice.sort()
        dice.reverse()
        yield dice

    def WhoWins(self, attack_dice, defence_dice):
        assert 1 <= len(attack_dice) <= 3 
        assert 1 <= len(defence_dice) <= 2 
        assert len(defence_dice) <= len(attack_dice)
        def_lose = att_lose = 0

        if ( attack_dice[ 0 ] > defence_dice[ 0 ] ):
            def_lose += 1
        else:
            att_lose += 1

        if ( len(defence_dice) <= 1 ):
            pass
        elif ( attack_dice[ 1 ] > defence_dice[ 1 ] ):
            def_lose += 1
        else:
            att_lose += 1
        return (att_lose, def_lose)

    def Effect(self, attack_dice, defence_dice):
        (att_lose, def_lose) = self.WhoWins(attack_dice, defence_dice)
        self.num_attacker_armies -= att_lose
        self.num_defender_armies -= def_lose
        if ( self.verbose ):
            print '... %s' % DESCRIPTION.get((att_lose, def_lose),
                    "attackers lose %d, defenders lose %d" % (
                            att_lose, def_lose))

    def MakeMove(self, num_attack_dice):
        assert 1 <= num_attack_dice <= 3
        assert num_attack_dice <= ( self.num_attacker_armies - 1 )
        if ( self.verbose ):
            print '... attacker rolls %d dice' % num_attack_dice

        for attack_dice in self.Produce(num_attack_dice):
            if ( self.verbose ):
                print '... attack roll is %s' % str(attack_dice)

            if ( num_attack_dice >= 2 ):
                # Can the defender choose how many dice to roll?
                has_choice = self.old_style_rules
            else:
                has_choice = False
                
            if ( has_choice ):
                # Apply rules from Risk FAQ 
                num_defence_dice = 1

                if ( self.num_defender_armies == 2 ):
                    # If the defender is down to exactly 2 armies,
                    # then he should roll 2 dice if the attacker's
                    # second die is 4 or less
                    if ( attack_dice[ 1 ] <= 4 ):
                        num_defence_dice = 2

                elif ( self.num_defender_armies == 4 ):
                    # If the defender has exactly 4 armies, he
                    # should roll 2 dice if the attacker's second
                    # die is 3 or less, or if the attacker's roll is
                    # 4-4.
                    if (( attack_dice[ 1 ] <= 3 )
                    or (( attack_dice[ 0 ] == 4 )
                    and ( attack_dice[ 1 ] == 4 ))):
                        num_defence_dice = 2

                elif ( attack_dice[ 1 ] <= 3 ):
                    # 3 armies, or 5 or more.
                    # 
                    # He should generally roll 2 dice if the attacker's
                    # second die is 3 or less.
                    num_defence_dice = 2

            else:
                # Roll as many as possible
                num_defence_dice = min(num_attack_dice, 2)
            
            # Limit to the number of armies available
            num_defence_dice = min(self.num_defender_armies, 
                        num_defence_dice)

            if ( self.verbose ):
                print '... defender rolls %d dice' % num_defence_dice

            for defence_dice in self.Produce(num_defence_dice):
                if ( self.verbose ):
                    print '... defence roll is %s' % str(defence_dice)
                self.Effect(attack_dice, defence_dice)
                
class TestEngine(Engine):
    def __init__(self):
        Engine.__init__(self, 10, 10, False)
        self.record = dict()

    def Produce(self, num_dice):
        assert 1 <= num_dice <= 3
        if ( num_dice == 1 ):
            recursive = [ [] ]
        else:
            recursive = self.Produce(num_dice - 1)

        for dice in recursive:
            for die in xrange(1, 7):
                dice2 = dice[:]
                dice2.append(die)
                dice2.sort()
                dice2.reverse()
                yield dice2

    def Test(self):
        print '\nFAQ example 3.2...\n'
        self.old_style_rules = False
        for num_attack_dice in xrange(1, 4):
            for num_defence_dice in xrange(1, 3):
                if ( num_defence_dice > num_attack_dice ):
                    continue
                   
                self.num_defender_armies = num_defence_dice
                print 'Attacker rolls %d dice, defender rolls %d dice' % (
                        num_attack_dice, num_defence_dice)
                self.TestRun(num_attack_dice)

        print '\nSmart algorithm from FAQ...\n'
        self.old_style_rules = True
        for num_defenders in xrange(1, 6):
            self.num_defender_armies = num_defenders
            for num_attack_dice in xrange(1, 4):
                print ( 'Attacker rolls %d dice, defender has %d armies,\n'
                    ' defender looks at attack roll before choosing.' %
                        (num_attack_dice, num_defenders))
                self.TestRun(num_attack_dice)

    def TestRun(self, num_attack_dice):
        self.record.clear()
        self.MakeMove(num_attack_dice)
        total = 0
        situation = []
        for ((att_lose, def_lose), count) in self.record.iteritems():
            total += count
            situation.append((count, att_lose, def_lose))

        situation.sort()
        situation.reverse()

        for (count, att_lose, def_lose) in situation:
            print '  att lose %d def lose %d: %1.2f%% (%d/%d)' % (
                    att_lose, def_lose, 
                    100.0 * float(count) / float(total),
                    count, total)
        print ''

    def Effect(self, attack_dice, defence_dice):
        (att_lose, def_lose) = self.WhoWins(attack_dice, defence_dice)
        key = (att_lose, def_lose)
        self.record.setdefault(key, 0)
        self.record[ key ] += 1

def Cmd(name):
    if ( tty == None ):
        # Platforms without tty/termios (e.g. Windows)
        print "%s Enter 'y' to confirm, anything else to cancel." % name
        line = sys.stdin.readline().strip().upper()
        return ( line.startswith('Y') )
    else:
        # Platforms with tty and termios (e.g. Linux, MacOS)
        print "%s Enter to confirm, any other key to cancel." % name
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        ch = '\0'
        try:
            tty.setraw(fd)
            ch = sys.stdin.read(1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
        return ( ch in "\n\r" )
    
def Num(title):
    while True:
        print '\r%s\n> ' % (title),
        line = sys.stdin.readline().strip()
        try:
            value = int(line)
        except ValueError:
            value = -1
        if ( 1 <= value ):
            return value

def Main():
    print ""
    print "Risk Assistant"
    while True:
        print ""
        print "NEW ATTACK"
        orig_num_attacker_armies = Num(
            "Please enter the number of attacking armies:")
        orig_num_defender_armies = Num(
            "Please enter the number of defending armies:")
        print ""
        engine = Engine(orig_num_attacker_armies, 
                        orig_num_defender_armies, True)
        num_attacks = 1
        attack = True
        while ( attack ):
            print ""
            print "[BATTLE %d] ATTACKERS %d   DEFENDERS %d" % (
                    num_attacks, engine.num_attacker_armies, 
                    engine.num_defender_armies)
            print ""

            attack = False
            if ( engine.num_defender_armies <= 0 ):
                print "The attacker has won! No defenders remain!"
            elif ( engine.num_attacker_armies < 2 ):
                print ( "The defense is successful! No further attacks "
                    "are possible." )
            else:
                if ( Cmd("Attack?") ):
                    num_attack_dice = min(3, engine.num_attacker_armies - 1)
                    engine.MakeMove(num_attack_dice)
                    num_attacks += 1
                    attack = True

        if ( num_attacks > 1 ):
            print "Results: "
            loss1 = orig_num_attacker_armies - engine.num_attacker_armies
            loss2 = orig_num_defender_armies - engine.num_defender_armies
            print "  In total, the attacker lost %d armies." % loss1
            print "  In total, the defender lost %d armies." % loss2 
            print "  %d attackers remain." % (engine.num_attacker_armies)
            print "  %d defenders remain." % (engine.num_defender_armies)
            max_loss = max(loss1, loss2)
            if (( max_loss >= 30 )
            and ( engine.num_defender_armies == 0 )):
                msg = random.randint(0, 3)
                if ( msg == 0 ):
                    print ( "  The death of one army is a tragedy, "
                        "the death of %d is just a statistic." % (max_loss) )
                elif ( msg == 1 ):
                    print ( "  The military/industrial complex thanks you "
                        "for your valuable assistance." )
                elif ( msg == 2 ):
                    print ( "  The Hague called, they want to talk to you "
                        "about some war crimes." )
                else:
                    print ( "  Now imagine how long that would have taken "
                        "using dice." )

    #te = TestEngine()
    #te.Test()

if ( __name__ == "__main__" ):
    try:
        Main()
    except KeyboardInterrupt:
        print ""
        print "Quit"

""" 
The Risk FAQ's "attacker rolls first" rule:
http://www.kent.ac.uk/IMS/personal/odl/riskfaq.htm

If the defender gets to see the attacker's roll before deciding, then 
he should generally roll 2 dice if the attacker's second die is 3 or 
less and 1 die if the attacker's second die is 4 or more. If the defender 
is down to exactly 2 armies, then he should roll 2 dice (assuming the 
rules permit him to) if the attacker's second die is 4 or less. This is 
because rolling one die and losing a single army forces the defender 
to roll only one die the next time, so he's better off trying to get 
lucky while he can still roll two dice. If the defender has exactly 
3 armies, however, he should still use the "second die is 3 or less" 
rule, rather than playing conservatively to avoid losing two armies. 
If the defender has exactly 4 armies, he should roll 2 dice if the 
attacker's second die is 3 or less, or if the attacker's roll is 
4-4. With 5 or more armies, always use the "3 or less" rule.
"""
