#!/usr/bin/python
# simulate coin tosses and check with chisq values from R

from random import randint
from math import exp, sqrt
import sys

# run n trials
def chisim(l, s, lim, n):
  x = [ 0 ] * l
  obs = [ 0 ] * l
  cnt = 0
  j = 0
  while j < n:
    x = cointoss(l, s)
    j += 1
    t = teststat(x)
    if j < 20: print x, ("%.4f" % t)
    for k in range(l): obs[k] += x[k]
    if (t >= lim): cnt += 1
  for k in range(l): obs[k] /= float(n)
  print "observed avg:", obs
  print "chisq test stat >=", lim, ":", cnt/float(n)

# toss a coin with l sides for s times
# return number of heads and tails (and ?)
def cointoss(l, s):
  x = [ 0 ] * l
  for j in xrange(s):
    i = randint(0, l-1)
    x[i] += 1
  return x

# test statistic for chi-square test
def teststat(x):
  n = len(x)
  e = sum(x)/float(n)
  t = 0.0
  for i in range(n):
    t += (x[i] - e)**2 / e
  return t

# l = sides of coin, s = tosses
# lim = compare value of test statistic, n = number of trials
def main():
  l = int(sys.argv[1])
  s = int(sys.argv[2])
  lim = float(sys.argv[3])
  n = int(sys.argv[4])
  x = [22, 40]
  print "check:", x, teststat(x)
  chisim(l, s, lim, n)

if __name__ == '__main__': main()