__author__ = "Maryam Bahrani"

""" a Module for plotting the ehrenfest distribution """
""" for help type into commandline: python plot.py --help
    
    sample usage:
    $ python plot.py 
    Number of balls to start with: 1000
    Number of swaps to do: 1000
    Number of iterations: 1000
"""

import plotly.offline as py
import plotly.graph_objs as go
import numpy as np
import click

def hist(n, s, m):
    data = []
    for i in range(m):
        a = n
        b = 0
        for j in range(s):
            x = np.random.randint(0,2)
            if a == 0 or (x == 0 and a != n):
                a = a + 1
                b = b - 1
            else:
                a = a - 1
                b = b + 1
        data.append(a)
    return data


@click.command()
@click.option('--num', prompt = 'Number of balls to start with',
              default = "1000",
              help = 'Number of balls in bin A at the start.')
@click.option('--stop', prompt = 'Number of swaps to do',
              default = "1000",
              help = 'Number of swaps to do.')
@click.option('--it', prompt = 'Number of iterations',
              default = "1000",
              help = 'Number of times to run the simulation.')
def run(num, stop, it):
    """ Enumerates labeled and unlabeled block graphs """
    x = hist(int(num), int(stop), int(it))
    data = [go.Histogram(x=x, histnorm='count')]
    py.plot(data)

if __name__ == '__main__':
  run()