#! /usr/bin/python
 
# Tools for randomness
# randint - pick a random integer in the given closed interval
# uniform - generate a (floating point) random number from the given interval
from random import randint, uniform
from math import sqrt
import pylab
 
def setupGraph(title, xlabel, ylabel, hold=False):
    pylab.title(title)
    pylab.xlabel(xlabel)
    pylab.ylabel(ylabel)
    pylab.grid(True)
    pylab.hold(hold)
 
def ideal_gas(
        N,                              # Number of particles
        totalEnergy,                    # total of demon and system energy
        steps,                          # number of simulation steps
        state=1,                        # Initial state
        visuals=True                    # plot histogram?
    ):
    v0 = sqrt(2.*totalEnergy/N)
    delta = v0/10
 
    # Initialize the system's state
    if state == 1:
        v = [v0]*N
        demonEnergy = 0                 # set initial demon energy
        systemEnergy = totalEnergy      # set initial system energy
    elif state == 2:
        v = [0]*N
        demonEnergy = totalEnergy       # set initial demon energy
        systemEnergy = 0                # set initial system energy
    else:                               # Bad state specified
        return 0                        # Really should raise ValueError
 
    Accum_systemEnergy = 0              #initialize our energy accumulators 
    demonEnergies = []
    norm = 1./steps
 
    demonUpdateInterval = steps/30      # Sample the demon 30 times
    demonAvgs = []
 
    for s in range(steps):
        for _ in range(N):
            which = randint(0,N-1)
            dv = uniform(-delta,delta)
 
            vnew = v[which]+dv           #compute our change in energy
            deltaEnergy = (vnew**2 - v[which]**2)/2 
 
            if deltaEnergy < demonEnergy:#if our demon has enough energy
                v[which] = vnew          #accept this move, update variables
                systemEnergy += deltaEnergy
                demonEnergy -= deltaEnergy
 
        if visuals:
            demonEnergies.append(demonEnergy)
            if (s+1)%demonUpdateInterval == 0:
                avg = sum(demonEnergies[(s+1) - 
                          demonUpdateInterval:])/demonUpdateInterval
                demonAvgs.append(avg)
 
                pylab.figure(4)
                pylab.plot(demonAvgs, 'r-o',
                           xdata=range(demonUpdateInterval,
                                       s+2,
                                       demonUpdateInterval))
                pylab.yscale('log', subsy=range(10))
                setupGraph(
                    title="Demon Energy Over Time",
                    xlabel="steps",
                    ylabel="Energy"
                )
 
        Accum_systemEnergy += systemEnergy*norm
 
 
 
    if visuals:
        pylab.figure(2)
        pylab.hist(v)
        setupGraph(
            title="Final particle velocity distribution",
            xlabel="V",
            ylabel="Frequency"
        )
 
        pylab.figure(3)
        pylab.hist(demonEnergies)
        setupGraph(
            title="Demon Energy Histogram",
            xlabel="Demon Energy",
            ylabel="Frequency"
        )
 
    return Accum_systemEnergy
 
if __name__ == "__main__":
    from math import exp
    import pylab
 
    state = int(raw_input('Select state [1, 2]: '))
    if state not in (1,2):
        state = 1
 
    Ns = []
    totalEnergies = []
    systemEnergies = []
    pylab.ion()
    setupGraph(
        title="N vs. Energy",
        xlabel="N",
        ylabel="Energy",
        hold=True
    )
    pylab.plot(Ns, totalEnergies, 'b-o', Ns, systemEnergies, 'g-o')
    pylab.legend(('Total Energy', 'System Energy'))
 
    totalEnergy = 500
    for N in range(50,501,50):
        systemEnergy = ideal_gas(N, totalEnergy, 3000, state, N==500)
 
        Ns.append(N)
        totalEnergies.append(totalEnergy)
        systemEnergies.append(systemEnergy)
 
        pylab.figure(1)
        pylab.plot(Ns, totalEnergies, 'b-o', Ns, systemEnergies, 'g-o')
    pylab.show()