#! /usr/bin/env python

from __future__ import print_function

# Make a topo plot of channel RMS.

import sys, os
import matplotlib
matplotlib.use('Agg')	# guard against no $DISPLAY
import matplotlib.pyplot as plt
import numpy as np
import pyctf
from pyctf.sensortopo import sensortopo
from pyctf.samiir import *

# Usage: RMStopo.py $ds setname [HP]
# The setname is used in the plot title, and the optional HP specifies
# the use of a high-pass filter.

hp = False
dsname = sys.argv[1]
setname = sys.argv[2]
if len(sys.argv) == 4:
    hp = True

# Open the dataset.

ds = pyctf.dsopen(dsname)
N = ds.getNumberOfSamples()
srate = ds.getSampleRate()

# Create a topo.

topo = sensortopo(ds)
topo.make_grid(200)

# The topo only uses the channels that are actually in the dataset.
# Get the channel indices.

names = topo.get_names()
nchan = len(names)
cidx = np.array(ds.clist2idx(names))
cidx -= ds.getFirstPrimary()

# Create filters. A wide band for the RMS, and a powerline notch.
# We also do it with the lowest frequencies removed.

if hp:
    filt = mkiir(2., 150., srate); Title = "Channel RMS, {}, 2-150 Hz".format(setname)
else:
    filt = mkiir(0., 150., srate); Title = "Channel RMS, {}, 0-150 Hz".format(setname)
#notch60 = mkfft(62., 58., srate, N) # slower but better resolution
#notch120 = mkfft(122., 118., srate, N)	
notch60 = mkiir(64., 56., srate)
notch120 = mkiir(124., 116., srate)

X = ds.getPriArray(0)               # only look at the first trial
Y = np.zeros((nchan, N))            # array to hold the filtered data

for i in cidx:
    Y[i] = X[i] * 1e15              # convert to fT
    Y[i] = dofilt(Y[i], filt)       # filter
    Y[i] = dofilt(Y[i], notch60)
    Y[i] = dofilt(Y[i], notch120)
    Y[i] -= Y[i].mean()             # remove the mean from each channel

# Compute the root mean square for each channel.

rms = np.sqrt((Y**2).mean(axis = 1))

# Plot the result on the topo.

im, ticks = topo.plot(rms, zrange = 'auto')
#im, ticks = topo.plot(rms, zrange = 'auto', cmap = plt.cm.coolwarm)
plt.text(.5, 1.03, Title, ha = 'center', weight = 'bold', size = 'large')
cax = plt.axes([.85, .15, .03, .65])
plt.colorbar(im, cax, format = '%g', ticks = ticks, extend = 'neither')
plt.savefig("rms.png", dpi = 150)

# Report the top few channel names

w = rms.argsort()
p = ds.getFirstPrimary()
for i in range(-1, -5, -1):
    j = w[i]
    print(ds.getChannelName(j + p), "%.1f" % rms[j])