import sys
import pymc as mc
from scipy.special import zeta

def _model(data, discrete=True, xmin=1.):

    alpha = mc.Uniform('alpha', 0.,10.) 
    #alpha = mc.Exponential('alpha', 1. / 1.5)
    #xmin = mc.Uniform('xmin', min(data), max(data))
    #xmin = mc.Exponential('xmin', 1.)
    #print xmin.value

    @mc.stochastic(observed=True)
    def custom_stochastic(value=data, alpha=alpha, xmin=xmin, discrete=discrete):
        value = value[value >= xmin]
        if discrete == True:
            return np.sum(np.log(value**-alpha / zeta(alpha,xmin)))
        else:
            return np.sum(np.log((alpha-1) * xmin**(alpha-1) * value**-alpha))

    return locals()

def analyze(data, discrete=True, xmin=1.):
    model = mc.MCMC(_model(data,discrete,xmin)) 
    model.sample(5000) 

    print
    print(model.stats()['alpha']['mean'])

    mc.Matplot.plot(model) 

import powerlaw

test = powerlaw.Power_Law(xmin = 1., parameters = [2.], discrete=True)
simulated = test.generate_random(5000)
analyze(simulated, discrete=True)

test = powerlaw.Power_Law(xmin = 1., parameters = [3.], discrete=False)
simulated = test.generate_random(5000)
analyze(simulated, discrete=False)

import urllib

data = list()
for line in urllib.urlopen("http://tuvalu.santafe.edu/~aaronc/powerlaws/data/words.txt"):
    data.append(line.strip())
analyze(data)

analyze(data,xmin=7)

data = list()
for line in urllib.urlopen("http://tuvalu.santafe.edu/~aaronc/powerlaws/data/fires.txt"):
    data.append(line.strip())
analyze(data, discrete=False)

analyze(data, discrete=False, xmin=6324.)