import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from plotstyle import *

if __name__ == "__main__":
    n = 21
    nn = 10000
    maxl = 5
    
    indices = np.arange(n)
    rng = np.random.RandomState(19281)
    x1 = rng.rand(n)
    rng = np.random.RandomState(35281)
    x2 = rng.rand(n)
    rng = np.random.RandomState(35281)
    x3 = rng.rand(n)

    x = rng.rand(nn)
    lags = np.arange(-maxl, maxl+1, 1)
    corrs = [np.corrcoef(x[maxl:-maxl-1-l], x[maxl+l:-maxl-1])[0, 1] for l in lags]

    fig = plt.figure(figsize=cm_size(figure_width, 1.8*figure_height))
    spec = gridspec.GridSpec(nrows=2, ncols=2,
                             left=0.11, bottom=0.12, right=0.98, top=0.97,
                             wspace=0.4, hspace=0.6)
    
    ax = fig.add_subplot(spec[0, 0])
    show_spines(ax, 'lb')
    ax.plot(indices, x1, c=colors['blue'], lw=1, zorder=10)
    ax.scatter(indices, x1, c=colors['blue'], edgecolor='white', s=50, zorder=20)
    ax.set_xlabel('Index')
    ax.set_ylabel('Random number')
    ax.set_xlim(-1.0, n+1.0)
    ax.set_ylim(-0.05, 1.05)
    
    ax = fig.add_subplot(spec[0, 1])
    show_spines(ax, 'lb')
    ax.plot(indices, x2, c=colors['blue'], lw=1, zorder=10)
    ax.scatter(indices, x2, c=colors['blue'], edgecolor='white', s=50, zorder=20)
    ax.set_xlabel('Index')
    ax.set_ylabel('Random number')
    ax.set_xlim(-1.0, n+1.0)
    ax.set_ylim(-0.05, 1.05)
    
    ax = fig.add_subplot(spec[1, 1])
    show_spines(ax, 'lb')
    ax.plot(indices, x3, c=colors['blue'], lw=1, zorder=10)
    ax.scatter(indices, x3, c=colors['blue'], edgecolor='white', s=50, zorder=20)
    ax.set_xlabel('Index')
    ax.set_ylabel('Random number')
    ax.set_xlim(-1.0, n+1.0)
    ax.set_ylim(-0.05, 1.05)
    
    ax = fig.add_subplot(spec[1, 0])
    show_spines(ax, 'lb')
    ax.plot(lags, corrs, c=colors['red'], lw=1, zorder=10)
    ax.scatter(lags, corrs, c=colors['red'], edgecolor='white', s=50, zorder=20)
    ax.set_xlabel('Lag')
    ax.set_ylabel('Serial correlation')
    ax.set_xlim(-maxl-0.5, maxl+0.5)
    ax.set_ylim(-0.05, 1.05)

    fig.savefig("randomnumbers.pdf")
    plt.close()