#!/usr/bin/env python

import itertools
import random
import sys

NUM_COINS = 12


class Scale(object):
    def __init__(self, debug=False):
        self.num_uses = 0
        self.debug = debug

    def use(self, side1, side2, debug=False):
        """
        returns 0 if equal, -1 if side 1, 1 if side 2
        """
        if self.debug or debug:
            print "%s vs %s" % (side1, side2)
        self.num_uses += 1
        side1_sum = sum([c.weight for c in side1])
        side2_sum = sum([c.weight for c in side2])
        val = 0
        if side1_sum < side2_sum:
            val = -1
        elif side2_sum < side1_sum:
            val = 1
        return val


class ScaleResult(object):
    def __init__(self, val):
        self.val = val

    @property
    def left_heavy(self):
        return self.val > 0

    @property
    def right_heavy(self):
        return self.val < 0

    @property
    def equal(self):
        return self.val == 0


class SmartScale(object):
    def __init__(self, scale):
        self.scale = scale

    @property
    def num_uses(self):
        return self.scale.num_uses

    def use(self, side1, side2):
        return ScaleResult(self.scale.use(side1, side2))


class Coin(object):
    def __init__(self, number):
        self.number = number
        self.weighted = False
        self.weight = 1

    def make_weighted(self, force_heavy=False, force_light=False):
        self.weighted = True
        if (random.random() > 0.5 or force_heavy) and not force_light:
            self.weight = 1.5
        else:
            self.weight = 0.5

    def __repr__(self):
        return "C%s (%s)" % (self.number, self.weight)


class Result(object):
    def __init__(self, num_weighings, correct):
        self.num_weighings = num_weighings
        self.correct = correct

    def __repr__(self):
        return "Result: [corr: %s, num: %s]" % (
            self.correct, self.num_weighings)


def solve_two_with_known(scale, coins, known_good):
    coin = coins[0]
    weighing = scale.use([coin], [known_good])
    if weighing.equal:
        coin = coins[1]
        weighing = scale.use([coin], [known_good])
    return coin, weighing.left_heavy


def solve_four(scale, coins):
    """
    Case where you only have four coins
    """
    weighing = scale.use(coins[0:1], coins[1:2])
    if weighing.equal:
        return solve_two_with_known(scale, coins[2:4], coins[0])
    else:
        return solve_two_with_known(scale, coins[0:2], coins[3])


def strategy1(scale, coins):
    """
    Split coins into 8,4, then narrow down.

    Returns weighted coin and whether it weighs more.
    """
    l_coins = coins[0:8]
    r_coins = coins[8:12]
    weighing = scale.use(l_coins[0:4], l_coins[4:8])
    if weighing.equal:
        suspect = r_coins
    else:
        weighing = scale.use(l_coins[0:2], l_coins[2:4])
        if weighing.equal:
            suspect = l_coins[4:8]
        else:
            suspect = l_coins[0:4]
    return solve_four(scale, suspect)


def strategy2(scale, coins, debug=False):
    """
    Jim's strat
    """
    pile1 = coins[0:4]
    pile2 = coins[4:8]
    spare_coins = coins[8:12]
    scale.scale.debug = debug
    w = scale.use(pile1, pile2)
    if w.equal:
        # then it's a spare coin problem
        if debug:
            print "Spare coin block"
        w2 = scale.use(spare_coins[0:3], pile1[0:3])
        if w2.equal:
            # then 12 is fake
            w3 = scale.use(spare_coins[3:4], pile1[0:1])
            return spare_coins[3], w3.left_heavy
        else:
            heavier = w2.left_heavy
            w3 = scale.use(spare_coins[0:1], spare_coins[1:2])
            if w3.equal:
                return spare_coins[2], heavier
            else:
                if heavier:
                    if w3.left_heavy:
                        return spare_coins[0], True
                    if w3.right_heavy:
                        return spare_coins[1], True
                else:
                    if w3.left_heavy:
                        return spare_coins[1], False
                    if w3.right_heavy:
                        return spare_coins[0], False
    elif w.left_heavy:
        heavy_pile1 = w.left_heavy
        w2 = scale.use(pile1[0:2] + [pile2[0]], pile1[2:4] + [pile2[2]])
        if w2.left_heavy:
            # implies 1 or 2 are fake and heavy or 7 is fake and light
            if debug:
                print "1 or 2 are fake/heavy or 7 is fake and light"
            w3 = scale.use([pile1[0]], [pile1[1]])
            if w3.equal:
                if heavy_pile1:
                    return pile2[2], not heavy_pile1
                else:
                    return pile2[0], not heavy_pile1
            else:
                # if left heavy, that one is the fake
                if w3.left_heavy:
                    return pile1[0], heavy_pile1
                else:
                    return pile1[1], heavy_pile1
        elif w2.right_heavy:
            # 125 < 347, 1 or 2 are fake and light or 7 is fake and heavy?
            w3 = scale.use([pile1[2]], [pile1[3]])  # 3 vs 4
            if w3.equal:
                # 3 = 4
                if heavy_pile1:
                    return pile2[0], w2.left_heavy
                else:
                    return pile2[2], w2.right_heavy
            else:
                if w3.left_heavy:
                    return pile1[2], heavy_pile1
                else:
                    return pile1[3], heavy_pile1
        elif w2.equal:
            # then either 6 or 8 are fake and light
            w3 = scale.use([pile2[1]], [spare_coins[0]])
            if w3.equal:
                # 8
                return pile2[3], not heavy_pile1
            else:
                # 6 is fake
                return pile2[1], not heavy_pile1
    else:
        # pile2 weighs more
        # 561 vs 783 ?
        w2 = scale.use(pile2[0:2] + [pile1[0]], pile2[2:4] + [pile1[2]])
        if w2.equal:
            # then either 2 or 4 are fake and light
            w3 = scale.use([pile1[1]], [spare_coins[0]])
            if w3.equal:
                # 4 is light
                return pile1[3], False
            else:
                return pile1[1], False
        elif w2.left_heavy:
            # then 5 6 are heavy or 3 is light? 
            w3 = scale.use([pile2[0]], [pile2[1]])
            if w3.equal:
                # 3 is light
                return pile1[2], False
            else:
                if w3.left_heavy:
                    return pile2[0], True
                else:
                    return pile2[1], True
        else:
            # then 7 8 are heavy or 1 is light
            w3 = scale.use([pile2[2]], [pile2[3]])
            if w3.equal:
                return pile1[0], False
            else:
                if w3.left_heavy:
                    return pile2[2], True
                else:
                    return pile2[3], True
            return pile1[0], False


def run_strat_loop(strategy, num_iter=1000):
    uses = []
    correctness = []
    for i in xrange(num_iter):
        result = run_once(strategy)
        uses.append(result.num_weighings)
        correctness.append(result.correct)
    return uses, correctness


def run_once(strategy, coins=None, debug=False):
    if debug:
        print "\n\n---- BEGIN RUN ----"
    scale = SmartScale(Scale())
    if not coins:
        coins = [Coin(x) for x in xrange(NUM_COINS)]
        if debug:
            print "Created coins 0-11"
        rand_idx = random.randint(0, NUM_COINS - 1)
        if debug:
            print rand_idx
        coin = coins[rand_idx]
        coin.make_weighted()
    else:
        coin = [c for c in coins if c.weight != 1][0]
    if debug:
        print "Made coin %s weighted: %s" % (coin.number, coin.weight)
        print "------------"
        print "Running strategy..."
    coin_found, weighs_more = strategy(scale, coins, debug=debug)
    if debug:
        print "Scale used: %s times" % scale.num_uses
        print "Coin found: %s" % coin_found.number
        print "Coin weighs %s than reference coin" % ("MORE" if weighs_more else "LESS")
    if ((weighs_more and coin.weight > 1)
            or (not weighs_more and coin.weight < 1)) and coin_found.number == coin.number:
        correct = True
    else:
        correct = False
    return Result(scale.num_uses, correct)


def crunch(num_trials):
    num_trials = int(num_trials)
    print "Running %s trials..." % num_trials
    weighings, corrects = run_strat_loop(strategy1, num_trials)
    print "average num. weighings: %s" % (int(sum(weighings)) / float(num_trials))
    print "correctness: %s" % (len([c for c in corrects if c is True]) / float(num_trials))


def edge_case_analyze(debug=False):
    if debug:
        debug = True
    all_results = []
    for x in xrange(2):
        if x == 0:
            kw = dict(force_heavy=True)
        else:
            kw = dict(force_light=True)
        results = {}
        for i in xrange(NUM_COINS):
            coins = [Coin(x) for x in xrange(NUM_COINS)]
            coins[i].make_weighted(**kw)
            results[i] = run_once(strategy2, coins=coins, debug=debug)
        all_results.append(results)
    import pprint
    pprint.pprint(all_results)


def test_once():
    result = run_once(strategy2, debug=True)
    print "Correct: %s" % result.correct


def main():
    m = {
        'crunch': crunch,
        'analyze': edge_case_analyze,
        'test': test_once
    }
    if len(sys.argv) > 1:
        arg = sys.argv[1]
        if len(sys.argv) > 2:
            arg2 = sys.argv[2]
            m[arg](arg2)
        else:
            m[arg]()
    else:
        m['test']()

if __name__ == "__main__":
    main()
