import random, urllib

## Identifiers/colors for Republican, Democrat, All
D, R, A = 0x0000FF, 0xFF0000, 0x660099

## Distribution of voters and candidates.

## Nate had a uniform distribution

Duniform   = [(i, D) for i in range(1, 60+1)]
Runiform   = [(i, R) for i in range(61, 100+1)]
Auniform   = Duniform + Runiform

## But you might also want candidates drawn from a Normal distribution

def Dnormal(N=60, s=15): return [(normal(30, s), D) for _ in range(N)]
def Rnormal(N=40, s=10): return [(normal(80, s), R) for _ in range(N)]
def Anormal(s1, s2): return Dnormal(s=s1) + Rnormal(s=s2)

## Or you could think of the voters as one big pool

def Anormal2(N=100, mu=50, s=20): return [DorR(normal(mu, s)) for _ in range(N)]
def Dnormal2(s=20): return [(x, p) for (x, p) in Anormal2(s=s) if p == D]
def Rnormal2(s=20): return [(x, p) for (x, p) in Anormal2(s=s) if p == R]

def DorR(x): return (x, (x<=60) and D or R)

## The structure for two-party and jungle primary election scenarios:

def twoparty(Ds, Rs):
    d = election([pick(Ds), pick(Ds)], Ds)
    r = election([pick(Rs), pick(Rs)], Rs)
    return election([d, r], Ds+Rs)

def jungle(dist, Ncandidates=4):
    candidates = [pick(dist) for _ in range(Ncandidates)]
    top2 = election(candidates, dist, nwinners=2)
    return the(election(top2, Auniform))

## Running an election and a simulation:

def election(candidates, voters, nwinners=1):
    votes = {}
    for v in voters:
        _, c = min((abs(c[0]-v[0]),c) for c in candidates)
        votes[c] = votes.get(c,0) + 1
    return the(top(nwinners, candidates, key=lambda c: votes.get(c,0)))

def simulate(scenario, title, S=10000):
    "Run the scenario S times, and return a chart of the results."
    winners = [scenario() for _ in range(S)]
    return barchart(histogram(winners, scale=S/500), title)
   
def histogram(winners, width=4, scale=1):
    "Given [(x, color)...], create bins of [(count, avg_color)...]"
    table = [ [] for i in range(1+100/width)]
    for (x, color) in winners:
        table[roundint(x/width)].append(color)
    return [(len(colors)/scale, blend(colors)) 
            for colors in table]

## Charting the results

def barchart(bins, title): 
    "Create a barchart from (count, color) bins."
    counts=','.join(str(count) for (count, _) in bins)
    colors = '|'.join('%s'%color for (_, color) in bins)
    return chart(title, '&cht=bvg&chbh=a,6,6&chd=t:%s&chco=%s' 
                 % (counts, colors))

def chart(title, options, border=1, width=400, height=285):
    title = urllib.quote_plus(title)
    return '<img border=%s width=%d height=%d src="http://chart.apis.google.com/chart?chtt=%s&chs=%dx%d%s">' % (
        border, width, height, title, width, height, options)

## Utilities

pick = random.choice

def normal(mu, s): return max(0, min(roundint(random.gauss(mu, s)), 100))

def top(n, seq, key=None): return list(sorted(seq, key=key, reverse=True))[:n]

def the(seq): 
    if len(seq)==1: return seq[0]  
    else: return seq

def roundint(x): return int(round(x))

def blend(colors): 
    if not colors: return '000000'
    R = sum(bool(c & 0xFF0000) for c in colors)
    B = sum(bool(c & 0x0000FF) for c in colors)
    return '%02x00%02x' % (255*R/(R+B), 255*B/(R+B))

## Create charts

print simulate(lambda: twoparty(Duniform, Runiform),
               "2 Party Primary; D: Uniform(1,60), R: Uniform(61, 100)")

print simulate(lambda: jungle(Auniform),
               "Jungle Primary; Uniform(1, 100)")

print simulate(lambda: twoparty(Dnormal(s=15), Rnormal(s=10)),
               "2 Party Primary; D: Normal(30, 15), R: Normal(80, 10)")

print simulate(lambda: jungle(Anormal(15, 10)),
               "Jungle Primary; D: Normal(30, 15), R: Normal(80, 10)")

print simulate(lambda: twoparty(Dnormal(s=30), Rnormal(s=20)),
               "2 Party Primary; D: Normal(30, 30), R: Normal(80, 20)")

print simulate(lambda: jungle(Anormal(30, 20)),
               "Jungle Primary; D: Normal(30, 30), R: Normal(80, 20)")

print barchart(histogram(Dnormal(30000)+Rnormal(20000), scale=100),
               "Voter distribution: Normal(30, 15), Normal(80, 20)")

print barchart(histogram(Anormal2(50000, s=20), scale=100),
               "Voter distribution: Normal(50, 20)")


print simulate(lambda: twoparty(Dnormal2(s=20), Rnormal2(s=20)),
               "2 Party Primary; single Normal(50, 20)")

print simulate(lambda: jungle(Anormal2(s=20)),
               "Jungle Primary; single Normal(50, 20)")