#!/usr/bin/env python3

"""
Selectively plots data from the *_snapspecs.dat files output by NDP.
"""

import matplotlib; matplotlib.use('Agg')

import argparse
import struct
import sys
import numpy as np
import subprocess
import time

from lsl.common import stations


class SnapData:
    colors = {'X': '#1f77b4', 'Y': '#ff7f0e'}
    
    def __init__(self, filename):
        with open(filename, 'rb') as f:
            nStands, nChans = struct.unpack('ll', f.read(16))
            self.specs = np.fromfile(
                f, count = nStands * 2 * nChans, dtype = np.float32
                ).reshape(nStands, 2, nChans)
        self.freqs = np.arange(nChans) * 196e6 / 8192
        self.specs = np.ma.array(self.specs, mask=np.where(self.specs == 0, 1, 0))
        self.stdIndices = {}
        lwa1 = stations.parse_ssmif('/home/op1/MCS/tp/SSMIF_CURRENT.txt')
        for i, a in enumerate(lwa1.antennas):
           if i % 2 == 0:
               self.stdIndices[a.stand.id] = i // 2
    
    def plotMany(self, stds, filename=None, expected=-1, actual=-1, currents=None):
        n = len(stds)
        nCols = int(np.ceil(np.sqrt(n)))
        nRows = int(np.ceil(n / float(nCols)))
        
        if filename is None:
            fig = plt.figure()
        else:
            fig = matplotlib.figure.Figure(figsize=(9,9))
        fig.subplots_adjust(wspace = 0, hspace = 0, left=0.12, bottom=0.12)
        if expected != -1 and actual != -1:
            diff = 1.0*(actual-expected)/expected
            if diff > 0:
                diff_message = " (up %.0f%%)" % (abs(diff)*100,)
            elif diff < 0:
                diff_message = " (down %.0f%%)" % (abs(diff)*100,)
            else:
                diff_message = ' (no change)'
            fig.suptitle('%s\nexpected %i bad; found %i%s' % (inFileName, expected, actual, diff_message))
        else:
            fig.suptitle('%s'% inFileName)
        axes = []
        
        if currents is None:
            currents = [-1]*self.specs.shape[0]*2
            
        x = self.freqs / 1e6
        
        xMin, xMax = 0, 98
        yMin, yMax = +9e99, -9e99
        
        for iPlot, std in enumerate(stds):
            leftCol = (iPlot % nCols == 0)
            bottomRow = (iPlot // nCols + 1 == nRows)
            ax = fig.add_subplot(nRows, nCols, iPlot + 1)
            
            i = self.stdIndices[std]
            for j, pol in enumerate(['X', 'Y']):
                y = 10 * np.log10(self.specs[i, j, :])
                ax.plot(x, y, color = self.colors[pol])
                if yMax < y.max(): yMax = y.max()
                if yMin > y.min(): yMin = y.min()
            
            if leftCol:
                if iPlot / nCols == (nRows - 1) / 2:
                    ax.set_ylabel('Power (dB)')

            else:
                ax.set_yticklabels([])
            if bottomRow:
                if iPlot % nCols == (nCols - 1) / 2:
                    ax.set_xlabel('Frequency (MHz)')
            else:
                ax.set_xticklabels([])
            
            axes.append(ax)
        
        # TODO: Come up with more appropriate limits
        #yMin, yMax = yMin, yMin + 60
        yMin, yMax = yMax - 80, yMax - 20
        feeGoodMin, feeGoodMax = 0.240, 0.260
        for ax, std in zip(axes, stds):
            ax.text(xMin * 0.05 + xMax * 0.95, yMin * 0.05 + yMax * 0.95,
                    '%d' % std, ha = 'right', va = 'top')
            for p in (0, 1):
                fee_cur = currents[2*(std-1)+p]
                
                if fee_cur >= 0.0:
                    if fee_cur >= feeGoodMin and fee_cur <= feeGoodMax:
                        fee_clr = 'black'
                    elif fee_cur > feeGoodMax:
                        fee_clr = 'blue'
                    else:
                        fee_clr = 'red'
                    ax.text(xMin + (xMax - xMin)*(0.25+0.5*p), yMin + (yMax - yMin)*0.16,
                            '%.0f' % (fee_cur*1000,), ha = 'center', va = 'top', color=fee_clr)
            ax.set_xlim(xMin, xMax)
            ax.set_ylim(yMin, yMax)
        
        if filename is None:
            plt.show()
        else:
            canvas = matplotlib.backends.backend_agg.FigureCanvasAgg(fig)
            canvas.print_figure(filename, dpi = 80)

class FEEData:
    def __init__(self, filename):
        self.filename = filename
        self._cur = None
        self._load()
        
    def _load(self):
        if self.filename is None:
            return False
            
        try:
            last_entry = subprocess.check_output(['tail', '-n', '1', self.filename])
            last_entry = last_entry.decode('ascii', errors='ignore')
            fields = last_entry.split(',')
            t_entry = float(fields[0])
            age = time.time() - t_entry
            if age > 600:
                print(f"WARNING: FEE current measurements are {age/60:.1f} min old, ignoring")
                return False
                
            self._cur = [float(v) for v in fields[2:]]
        except subprocess.CalledProcessError:
            return False
            
        return True
        
    def get(self):
        return self._cur

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description =
        'Plots the spectra of the specified STD(s) for the given SNAP2 auto-'
        'correlation spectra file from NDP.  Blue shows X pol and orange shows '
        'Y pol.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
        )
    parser.add_argument('inputname', type=str, 
                        help='filename to process')
    parser.add_argument('stand', type=str, nargs='+',
                        help='stand(s) to plot')
    parser.add_argument('-o', '--outfile', type=str, 
                        help='write output to FILE; default is to plot to screen')
    parser.add_argument('-e', '--expected-bad', type=int, default=-1,
                        help='expected number of bad dipoles')
    parser.add_argument('-a', '--actual-bad', type=int, default=-1,
                        help='actual number of bad dipoles')
    parser.add_argument('-f', '--fee-currents', type=str,
                        help='FEE currents file from ARX to use for annotation')
    args = parser.parse_args()
    
    inFileName = args.inputname
    stds = [int(arg, 10) for arg in args.stand]
    
    if args.outfile is not None:
        import matplotlib.figure
        import matplotlib.backends.backend_agg
    else:
        import matplotlib.pyplot as plt
    
    snapdata = SnapData(inFileName)
    feedata = FEEData(args.fee_currents)
    snapdata.plotMany(stds, filename=args.outfile, expected=args.expected_bad, actual=args.actual_bad, currents=feedata.get())