#!/usr/bin/env python
from __future__ import division
import numpy as np
import matplotlib.pyplot as pp

"""
This script processes and plots channel data from the Panoptes visual
field mapping classroom experiment. Event times are automatically
drawn onto the light sensor data, with corresponding location indices
written for convenience. It can be run on any computer with
the a Python environment, and the numerics/graphics libraries indicated
in the lines above. Alternatively, it can be run through a web interface
called Python Anywhere. Instructions for that process are found in a
separate document.

When this script runs, it looks for files matching "light_log*.txt" and
"timestamp*.txt". For example, the data files for three groups can be
re-named:

light_log1.txt
light_log2.txt
light_log3.txt
timestamp1.txt
timestamp2.txt
timestamp3.txt

Files will be processed in pairs, and plots will be made for each group.
Within group, the 4 sensor channels will be mixed up randomly to enhance
the challenge of reconstructing the origin.

The grid size for the light mask can be modified by changing the
"default_grid" parameters just below. If you choose to make more than 3
repetitions in the experiment, the number of passes across the grid
is automatically detected.

For the benefit of the teacher, a computer generated map of the "receptive
fields" are plotted and saved as well. This map is very roughly approximated,
and its quality depends on the reliability of the event timestamps.

Finally, the script gives a warning if the actual number of timestamps does
not match the expected number of events.
"""

# The grid size on the stencil is 6 x 5 (width by height).
# Feel free to change it for more resolution and/or more coverage!
default_grid = (6, 5)
#default_grid = (8, 5)

# A4 page size in landscape (wid, ht)
psize = (11.69, 8.27)

class DataError(Exception):
    pass

def z_score(arr):
    ashape = arr.shape
    zarr = arr.ravel()
    zarr -= arr.mean()
    zarr /= arr.std()
    return zarr.reshape(ashape)

# Note: this method always plots the record across 4 rows in landscape layout
def make_page_fig2(x, ticks, thresh, ttl='', ncond=40):
    f = pp.figure(figsize=psize)
    nx = len(x)
    nx_row = nx//4
    # we'll truncate any remainder points, but don't miss the last
    # tick in the event that it falls within those points
    if ticks[-1] >= 4*nx_row:
        ticks[-1] = 4*nx_row-1

    xax = np.arange(1,len(x)+1)
    amx = x.max(); amn = x.min(); rng = amx - amn
    lims = (0.5*(amx+amn) - 0.55*rng, 0.5*(amx+amn) + 0.75*rng)

    cond_cnt = np.arange(len(ticks)) % ncond + 1
    for n in xrange(4):
        ax = f.add_subplot(4,1,n+1)
        ax.plot(xax[n*nx_row:(n+1)*nx_row], x[n*nx_row:(n+1)*nx_row])
        tmask =  (ticks >= n*nx_row) & (ticks < (n+1)*nx_row)
        sub_ticks = ticks[ tmask ]
        sub_cond_cnt = cond_cnt[ tmask ]
        for m, t in zip(sub_cond_cnt, sub_ticks):
            ax.axvline(t, color='r', linestyle=':')
            ax.text(t, amx*1.1, '%d'%m, fontsize=10)
        ax.axhline(thresh, color='k', linestyle='--')
        ax.set_ylim(lims)
        if ttl and n==0:
            ax.set_title(ttl)
    ax.set_xlabel('time', fontsize='large')
    f.text(0.01, 0.5, 'light intensity', rotation='vertical', fontsize='large')
    #f.tight_layout()
    f.subplots_adjust(
        left=0.06, bottom=0.08, top=0.95, right=0.97,
        wspace=0.2, hspace=0.224
        )
    return f

# Note: this method always plots one row for each pass through the grid
def make_page_fig(x, ticks, thresh, ttl='', ncond=40):
    nrep = len(ticks) // ncond

    ticks = ticks.reshape(nrep, ncond)
    breaks = 0.5 * (ticks[:-1,-1] + ticks[1:,0])
    breaks = np.r_[0, breaks, len(x)]

    xax = np.arange(1,len(x)+1)
    amx = x.max(); amn = x.min(); rng = amx - amn
    lims = (0.5*(amx+amn) - 0.55*rng, 0.5*(amx+amn) + 1.2*rng)

    f = pp.figure(figsize=psize)
    for n in xrange(nrep):
        ax = f.add_subplot(nrep, 1, n+1)
        ax.plot(xax[breaks[n]:breaks[n+1]], x[breaks[n]:breaks[n+1]])
        for m, t in enumerate(ticks[n]):
            ax.axvline(t, color='r', linestyle=':')
            ax.text(t, amx*1.1, '%d'%(m+1), fontsize=10)
        ax.axhline(thresh, color='k', linestyle='--')
        ax.set_ylim(lims)
        ax.set_xlim( (breaks[n], breaks[n+1]) )
        if ttl and n==0:
            ax.set_title(ttl)
    ax.set_xlabel('time', fontsize='large')
    f.text(0.01, 0.5, 'light intensity', rotation='vertical', fontsize='large')
    #f.tight_layout()
    f.subplots_adjust(
        left=0.06, bottom=0.08, top=0.95, right=0.97,
        wspace=0.2, hspace=0.224
        )
    return f

def proc_data(lsamp, ticks, gridsize=default_grid):
    # n_rep is the number of passes through the grid (default 3)
    # grid size is (width, height)
    # grid size is (6,5) by default, but can be varied in each classroom

    g_wd, g_ht = gridsize

    arr = np.loadtxt(lsamp)
    arr = arr[0:4*int(len(arr)/4)]
    arr = arr.reshape(-1, 4)

    stdev = np.std(arr, axis=0)

    # Response threshold is set here (for each channel)
    # A value that is 2.5x stdev past the mean has only about
    # %0.6 chance of spontaneously occurring from the background distribution.

    #thresh = arr.mean(axis=0) + 2.5*stdev
    thresh = arr.mean(axis=0) + 1.5*stdev

    # Load timestamps
    tx = np.loadtxt(ticks)

    # this should divide without remainder,
    # otherwise there are missing timestamps
    n_rep = len(tx) // (g_wd*g_ht)
    if len(tx) != g_wd*g_ht*n_rep:
        raise DataError

    # randomize the channels handed out to each group member
    rnd_chans = np.random.permutation(range(4))

    grp_figs = list()
    for g, n in enumerate(rnd_chans):
        f = make_page_fig(arr[:,n], tx, thresh[n],
                          ncond=g_wd*g_ht, ttl='Scientist %d Data'%(g+1))
        grp_figs.append(f)

    # This part makes a rough estimate of the Panoptes receptive
    # field for use by the instructor. It is simply the sum of responses
    # between events. In other words, the score for each cell is
    # the sum of all samples after the corresponding condition, but
    # before the next condition
    arr = arr - arr.mean(axis=0)
    cond_idx = np.arange(g_ht*g_wd).reshape(g_ht, g_wd)
    # every even row is reversed
    for n in xrange(1,g_ht,2):
        idx = cond_idx[n,::-1]
        cond_idx[n,:] = idx

    csum = np.cumsum(arr, axis=0)
    tx = tx.astype('i')
    tx_nxt = np.r_[tx[1:], arr.shape[0]-1]
    scores = csum[tx_nxt,:] - csum[tx,:]
    scores.shape = (n_rep, g_wd*g_ht, 4)
    scores = np.sum(scores, axis=0)
    scores = scores[cond_idx.ravel(),:]
    scores = z_score(scores)
    scores.shape = (g_ht, g_wd, 4)
    f = pp.figure(figsize=(8,8))

    for n in xrange(2):
        for m in xrange(2):
            ax = f.add_subplot(2, 2, n*2+m+1)
            ax.imshow(scores[:,::-1,n*2+m], interpolation='nearest',
                      origin='upper', extent=[0.5, g_wd+0.5, g_ht+0.5, 0.5])
            ax.set_title('Channel %d'%(n*2+m+1))
    return grp_figs, f, scores, arr


if __name__=='__main__':
    import sys, glob

    lsamps = glob.glob('light_log*.txt')
    ticks = glob.glob('timestamp*.txt')

    for n, data in enumerate(zip(lsamps, ticks)):
        print 'Processing group %d'%(n+1)
        g_samps, g_ticks = data
        try:
            gfigs, ifig, scores, arr = proc_data(g_samps, g_ticks)
        except DataError:
            print('There are not enough ticks for Group %d!!'%(n+1))
            continue

        for s in xrange(4):
            fname = 'grp_%d_plot_sci_%d.pdf'%(n+1, s+1)
            gfigs[s].savefig(fname)
            print 'saved figure: %s'%fname,

        print ''
        ifig.savefig('grp_%d_instructor_plot.pdf'%(n+1))

    ## Uncomment this line to show the plots (e.g. if you are running Python
    ## on a local machine, rather than on the web)
    # pp.show()