#!/usr/bin/env python
"""Convolutional Neural Network Training Functions

Functions for building and training a (UNET) Convolutional Neural Network on
images of the Mars and binary ring targets.
"""
from __future__ import absolute_import, division, print_function

import numpy as np
#import pandas as pd
import h5py

from keras.models import load_model
from keras.callbacks import ModelCheckpoint, TensorBoard, Callback
#from deepmars2.YNET.model import build_YNET, get_weight, update_weight_callback
from deepmars2.ResUNET.model import ResUNET
import deepmars2.features.template_match_target as tmt
import click
import logging
import os
from tqdm import tqdm
import time
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
from imgaug import augmenters as iaa
from joblib import Parallel, delayed
import config as cfg
#import pickle

# Reduce Tensorflow verbosity
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# allow free allocation of gpu memory
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.log_device_placement = True
sess = tf.Session(config=config)
set_session(sess)

@click.group()
def dl():
    log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    logging.basicConfig(level=logging.INFO, format=log_fmt)

def generator_with_replacement(data, target, image_type, batch_size=10):
    augmentation = iaa.Sequential([
            iaa.Fliplr(0.5),
            iaa.OneOf([iaa.Affine(rotate=0), iaa.Affine(rotate=90),
                       iaa.Affine(rotate=180), iaa.Affine(rotate=270)])])
    max_index = data.shape[0]

    while True:
        indices = [np.random.randint(max_index) for _ in range(batch_size)]
        d, t = data[indices].copy(), target[indices].copy()

        DEM = d[:,0,...]
        IR = d[:,1,...]
        DEM = np.squeeze(DEM.transpose([3,1,2,0]))
        IR = np.squeeze(IR.transpose([3,1,2,0]))
        t = np.squeeze(t.transpose([3,1,2,0]))

        det = augmentation.to_deterministic()

        DEM = det.augment_image(DEM)
        IR = det.augment_image(IR)
        t = det.augment_image(t)

        DEM = DEM.transpose([2,0,1])[..., np.newaxis]
        IR = IR.transpose([2,0,1])[..., np.newaxis]
        t = t.transpose([2,0,1])[..., np.newaxis]

        if image_type not in ['IR', 'DEM']:
            raise ValueError('image_type must be one of IR or DEM')
        if image_type == 'IR':
            yield (IR, t)
        if image_type == 'DEM':
            yield (DEM, t)

def t2c(pred, csv, i):
    return np.hstack([i,
                      tmt.template_match_t2c(pred, csv)][1:])

def diagnostic(res, beta):
    """Calculate the metrics from the predictions compared to the CSV.

    Parameters
    ------------
    res: list of results containing:
        image number, number of matched, number of existing craters, number of
        detected craters, maximum radius detected, mean error in longitude,
        mean error in latitude, mean error in radius, fraction of duplicates
        in detections.
    beta : int
        Beta value when calculating F-beta score.

    Returns
    -------
    dictionary : metrics stored in a dictionary
    """

    counter, N_match, N_csv, N_detect,\
        mrad, err_lo, err_la, err_r, frac_duplicates = np.array(res).T
    try:
        positives = N_detect.sum()
        true_positives = N_match.sum()
        ground_truth = N_csv.sum()
        precision = true_positives / positives
        recall = true_positives / ground_truth
        fscore = (1 + beta**2) * (recall * precision) / \
                 (precision * beta**2 + recall)
    except:
        # division by zero occured, so no craters were matched
        precision = 0
        recall = 0
        fscore = 0

    return dict(precision=precision, recall=recall, fscore=fscore)

def get_metrics(data, craters_images, model, preds, name,
                dim=256, beta=1, offset=0, rmv_oor_csvs=0):
    """Function that prints pertinent metrics at the end of each epoch.

    Parameters
    ----------
    data : hdf5
        Input images.
    craters : hdf5
        Pandas arrays of human-counted crater data.
    model : keras model object
        Keras model
    beta : int, optional
        Beta value when calculating F-beta score. Defaults to 1.
    """

    # Either a model is provided to make predictions or the predictions have
    # already been made
    assert (model is not None) or (preds is not None)

    X = data[0]
    craters, images = craters_images
    # Get csvs of human-counted craters
    csvs = []
    cutrad = 0.8
    n_csvs = len(X)
#    diam = 'Diameter (pix)'

    for i in range(len(X)):
        imname = images[i]  # name = 'img_{0:05d}'.format(i)
        found = False
        for crat in craters:
            if imname in crat:
                csv = crat[imname]
                found = True
        if not found:
            print('not found: ', imname)
            csvs.append([-2])
            continue
        # Remove small/large/half craters
        csv = csv[csv['Diameter (pix)'] < 2 * cfg.maxrad_]
        csv = csv[csv['Diameter (pix)'] > 2 * cfg.minrad_]
        csv = csv[csv['x (pix)'] + cutrad * csv['Diameter (pix)'] / 2 <= dim]
        csv = csv[csv['y (pix)'] + cutrad * csv['Diameter (pix)'] / 2 <= dim]
        csv = csv[csv['x (pix)'] - cutrad * csv['Diameter (pix)'] / 2 > 0]
        csv = csv[csv['y (pix)'] - cutrad * csv['Diameter (pix)'] / 2 > 0]
        if len(csv) < 3:    # Exclude csvs with few craters
            csvs.append([-1])
        else:
            csv_coords = np.asarray((csv['x (pix)'], csv['y (pix)'], csv['Diameter (pix)'] / 2)).T
            csvs.append(csv_coords)

    if model is not None:
        # A model has been provided so we must make predictions
        X = X.transpose([1,0,2,3,4])
        preds = model.predict([X[0], X[1]], verbose=1, batch_size=5)
        # Save predictions
        h5f = h5py.File('predictions.hdf5', 'w')
        h5f.create_dataset(name, data=preds)
        print('Successfully generated and saved model predictions.')

    # preds contains a large number of predictions,
    # so we run the template code in parallel.
    res = Parallel(n_jobs=16,
                   verbose=1)(delayed(t2c)(preds[i][:,:,0], csvs[i], i)
                              for i in range(n_csvs) if len(csvs[i]) >= 3)

    if len(res) == 0:
        print('No valid results: ', res)
        return None

    # At this point we've processed the predictions with the template matching
    # algorithm, now calculate the metrics from the data.

    return diagnostic(res, beta)

def test_model(Data, Craters, MP):

    if MP['model'] is not None:
        model = load_model(MP['model'])
        preds = None
    else:
        # no model specified so load predictions file
        h5f = h5py.File(MP['predictions'], 'r')
        preds = h5f[MP['test_dataset']][:]
        model = None

    diag = get_metrics(Data[MP['test_dataset']],
                       Craters[MP['test_dataset']],
                       model,
                       preds,
                       MP['test_dataset'])

    return diag

def train_and_test_model(Data, Craters, MP):
    """Function that trains, tests and saves the model, printing out metrics
    after each model.

    Parameters
    ----------
    Data : dict
        Inputs and Target Moon data.
    Craters : dict
        Human-counted crater data.
    MP : dict
        Contains all relevant parameters.
    i_MP : int
        Iteration number (when iterating over hypers).
    """

    # Parameters
    nb_epoch = MP['epochs']
    bs = MP['bs']
    learn_rate = MP['lr'][0]
    image_type = MP['image_type']

    # Build model
    if MP['model'] is not None:
        model = load_model(MP['model'])
    else:
#        model = build_YNET(4, 2, dim, FL, n_filters, drop, learn_rate,
#                           lmbda=lmbda, activation_function='relu')
        # taken from Zhang & Liu 2017
        n_kernels = 64
        kernel_size = 3
        depth = 3
        learn_rate = 0.0001
        dropout = 0.5
        #output_length = 2*(cfg.maxrad_ - cfg.minrad_)
        output_length = 1
        #model = ResUNET(n_kernels, kernel_size, depth, learn_rate, dropout)
        model = ResUNET(n_kernels, kernel_size, depth, learn_rate, dropout,
                        output_length=output_length)

    # Callbacks
    now = time.strftime('%c')
    n_samples = 5000
    save_folder = os.path.join(cfg.root_dir, 'ResUNET/models', now)
    os.mkdir(save_folder)
    save_name = save_folder + '/{epoch:02d}-{val_loss:.2f}.hdf5'
    save_model = ModelCheckpoint(os.path.join(cfg.root_dir, save_name),
                                 save_best_only=True)
    log_dir=os.path.join(cfg.root_dir, 'ResUNET/logs', now)
    tensorboard = TensorBoard(log_dir, batch_size=bs)
    tbi_test = TensorBoardImage(log_dir, [Data['dev'][0], Data['dev'][1]],
                                image_type)

    model.fit_generator(generator_with_replacement(Data['train'][0],
                                                   Data['train'][1],
                                                   image_type,
                                                   batch_size=bs),
                        steps_per_epoch=n_samples // bs,
                        verbose=1,
                        validation_data=generator_with_replacement(Data['dev'][0],
                                                                   Data['dev'][1],
                                                                   image_type,
                                                                   batch_size=bs),
                        validation_steps=300,
                        callbacks=[save_model, tensorboard, tbi_test],
                        epochs=nb_epoch)

def make_image(tensor):
    """
    Convert an numpy representation image to Image protobuf.
    Copied from https://github.com/lanpa/tensorboard-pytorch/
    """
    from PIL import Image
    height, width = tensor.shape
    tensor = (255*tensor).astype('uint8')
    image = Image.fromarray(tensor)
    import io
    output = io.BytesIO()
    image.save(output, format='PNG')
    image_string = output.getvalue()
    output.close()
    return tf.Summary.Image(height=height,
                         width=width,
                         colorspace=1,
                         encoded_image_string=image_string)

class TensorBoardImage(Callback):
    def __init__(self, log_dir, data, image_type):
        super().__init__()
        self.log_dir = log_dir
        self.data = data
        if image_type not in ['IR', 'DEM']:
            raise ValueError('image_type must be one of IR or DEM')
        self.image_type = image_type

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        data = self.data # self.validation_data
        input_IR = data[0][:,1]
        input_DEM = data[0][:,0]
        target_masks = data[1]

        print('gerenating validation prediction images')
        if self.image_type == 'IR':
            preds = self.model.predict(input_IR, batch_size=1)
        if self.image_type == 'DEM':
            preds = self.model.predict(input_DEM, batch_size=1)

        images = []
        for i in tqdm(range(10)):
            img = np.vstack([np.hstack([input_DEM[i,:,:,0],
                                        input_IR[i,:,:,0]]),
                             np.hstack([target_masks[i,:,:,0],
                                        preds[i,:,:,0]])
                            ])

            images.append(make_image(img))
        summary = tf.Summary(value=[
                tf.Summary.Value(tag='Image {:02d}'.format(i), image=images[i])
                for i in range(len(images))] +
                [tf.Summary.Value(tag='loss', simple_value=logs['loss']),
                 tf.Summary.Value(tag='val_loss', simple_value=logs['val_loss']),
#                 tf.Summary.Value(tag='weight', simple_value=get_weight())
                 ])
        writer = tf.summary.FileWriter(self.log_dir)
        writer.add_summary(summary, epoch)
        writer.close()

        return

def get_models(MP):
    """Top-level function that loads data files and calls train_and_test_model.

    Parameters
    ----------
    MP : dict
        Model Parameters.
    """
    dir = MP['dir']
    n_train, n_dev, n_test = MP['n_train'], MP['n_dev'], MP['n_test']
    prefix = MP['prefix']

    # Load data
    def load_files(numbers, test, this_dataset):
        res0 = []
        res1 = []
        files = []
        craters = []
        images = []
        npic = 0
        if not test or (test and this_dataset):
            for n in tqdm(numbers):
                files.append(h5py.File(os.path.join(
                    dir, prefix + '_images_{0:05d}.hdf5'.format(n)), 'r'))
                images.extend(['img_{0:05d}'.format(a)
                              for a in np.arange(n, n + 1000)])
                #res0.append(files[-1]['input_DEM'][:].astype('float32'))

                res0.append([files[-1]['input_DEM'][:].astype('float32'),
                             files[-1]['input_IR'][:].astype('float32')])
                npic = npic + len(res0[-1])
                res1.append(files[-1]['target_masks'][:].astype('float32'))
                files[-1].close()
#                craters.append(pd.HDFStore(os.path.join(
#                    dir, prefix + '_craters_{0:05d}.hdf5'.format(n)), 'r'))
            res0 = np.vstack(res0).transpose([1,0,2,3])
            res1 = np.vstack(res1)
        return files, res0, res1, npic, craters, images

    train_files,\
        train0,\
        train1,\
        Ntrain,\
        train_craters,\
        train_images = load_files(MP['train_indices'],
                                  MP['test'],
                                  MP['test_dataset'] == 'train')
    print(Ntrain, n_train)

    dev_files,\
        dev0,\
        dev1,\
        Ndev,\
        dev_craters,\
        dev_images = load_files(MP['dev_indices'],
                                MP['test'],
                                MP['test_dataset'] == 'dev')
    print(Ndev, n_dev)

    test_files,\
        test0,\
        test1,\
        Ntest,\
        test_craters,\
        test_images = load_files(MP['test_indices'],
                                 MP['test'],
                                 MP['test_dataset'] == 'test')
    print(Ntest, n_test)

    Data = {
        'train': [train0, train1],
        'dev': [dev0, dev1],
        'test': [test0[:n_test], test1[:n_test]]
        }

    def preprocess(Data):
        for key in Data:
            if len(Data[key][0]) == 0:
                continue
            for i in range(len(Data[key])):
                newdim = list(Data[key][i].shape) + [1]
                Data[key][i] = Data[key][i].reshape(*newdim)

    preprocess(Data)

    # Load ground-truth craters
    Craters = {
        'train': [train_craters, train_images],
        'dev': [dev_craters, dev_images],
        'test': [test_craters, test_images]
    }

#    # Iterate over parameters
    if MP['test']:
#        minrad_ = 3
#        maxrad_ = 60
#        longlat_thresh2_ = 1.8
#        rad_thresh_ = 1.0
#        template_thresh_ = 0.5
#        target_thresh_ = 0.1

        diag = test_model(Data, Craters, MP)
        print('Precision: {}'.format(diag['precision']))
        print('Recall: {}'.format(diag['recall']))
        print('F1: {}'.format(diag['fscore']))
    else:
        train_and_test_model(Data, Craters, MP)

@dl.command()
@click.option('--test', is_flag=True, default=False)
@click.option('--test_dataset', default='ran3')
@click.option('--model', default=None)
@click.option('--predictions', default=None)
@click.option('--prefix', default='ran3')
@click.option('--image_type', default=None)
def train_model(test, test_dataset, model, predictions, prefix, image_type):
    """Run Convolutional Neural Network Training

    Execute the training of a (UNET) Convolutional Neural Network on
    images of the Moon and binary ring targets.
    """

    # Model Parameters
    MP = {}

    # Directory of train/dev/test image and crater hdf5 files.
    MP['dir'] = os.path.join(cfg.root_dir, 'data/processed/')

    # Image width/height, assuming square images.
    MP['dim'] = 256

    # Batch size: smaller values = less memory, less accurate gradient estimate
    MP['bs'] = 10

    # Number of training epochs.
    MP['epochs'] = 800

    # Number of samples per epoch
    MP['per_epoch'] = 1000

    MP['train_indices'] = list(np.arange(0, 40000, 1000))
    MP['dev_indices'] = [41000, 42000, 43000, 44000, 45000]
    MP['test_indices'] = [46000, 47000, 48000, 49000]

    MP['n_train'] = len(MP['train_indices']) * 1000
    MP['n_dev'] = len(MP['dev_indices']) * 1000
    MP['n_test'] = len(MP['test_indices']) * 1000

    # initial model filename
    MP['model'] = model

    # initial predictions filename
    MP['predictions'] = predictions

    # testing only
    MP['test'] = test
    MP['test_dataset'] = test_dataset

    MP['filter_length'] = [3]           # Filter length
    MP['lr'] = [0.0001]                 # Learning rate
    MP['n_filters'] = [112]             # Number of filters
    MP['init'] = ['he_normal']          # Weight initialization
    MP['lambda'] = [1e-6]               # Weight regularization
    MP['dropout'] = [0.15]              # Dropout fraction

    MP['prefix'] = prefix

    if image_type not in ['IR', 'DEM']:
        raise ValueError('image_type must be one of IR or DEM')

    MP['image_type'] = image_type

    print(MP)
    get_models(MP)

if __name__ == '__main__':
    dl()


import numpy as np
import matplotlib.pylab as plt
import warnings
import os
from pathlib import Path
from ephysiopy.common.ephys_generic import PosCalcsGeneric
from ephysiopy.openephys2py.OESettings import Settings, OEStructure
from ephysiopy.visualise.plotting import FigureMaker

def fileExists(pname, fname) -> bool:
    return os.path.exists(os.path.join(pname, fname))

def fileContainsString(pname: str, searchStr: str) -> bool:
    if os.path.exists(pname):
        with open(pname, 'r') as f:
            strs = f.read()
        lines = strs.split('\n')
        for line in lines:
            if searchStr in line:
                return True
            else:
                return False
    else:
        return False

class KiloSortSession(object):
    """
    Loads and processes data from a Kilosort session.

    A kilosort session results in a load of .npy files, a .csv or .tsv file.
    The .npy files contain things like spike times, cluster indices and so on.
    Importantly	the .csv (or .tsv) file contains the cluster identities of
    the SAVED part of the phy template-gui (ie when you click "Save" from the
    Clustering menu): this file consists of a header ('cluster_id' and 'group')
    where 'cluster_id' is obvious (relates to identity in spk_clusters.npy),
    the 'group' is a string that contains things like 'noise' or 'unsorted' or
    whatever as the phy user can define their own labels.

    Parameters
    ----------
    fname_root : str
        The top-level directory. If the Kilosort session was run directly on
        data from an openephys recording session then fname_root is typically
        in form of YYYY-MM-DD_HH-MM-SS
    """
    def __init__(self, fname_root):
        """
        Walk through the path to find the location of the files in case this
        has been called in another way i.e. binary format a la Neuropixels
        """
        self.fname_root = fname_root
        import os
        for d, c, f in os.walk(fname_root):
            for ff in f:
                if '.' not in c:  # ignore hidden directories
                    if 'spike_times.npy' in ff:
                        self.fname_root = d
        self.cluster_id = None
        self.spk_clusters = None
        self.spk_times = None
        self.good_clusters = []

    def load(self):
        """
        Load all the relevant files

        There is a distinction between clusters assigned during the automatic
        spike sorting process (here KiloSort2) and the manually curated
        distillation of the automatic process conducted by the user with
        a program such as phy.

        * The file cluster_KSLabel.tsv is output from KiloSort.
            All this information is also contained in the cluster_info.tsv
            file! Not sure about the .csv version (from original KiloSort?)
        * The files cluster_group.tsv or cluster_groups.csv contain
            "group labels" from phy ('good', 'MUA', 'noise' etc).
            One of these (cluster_groups.csv or cluster_group.tsv)
            is from kilosort and the other from kilosort2
        """
        import os
        import pandas as pd
        dtype = {'names': ('cluster_id', 'group'), 'formats': ('i4', 'S10')}
        # One of these (cluster_groups.csv or cluster_group.tsv) is from
        # kilosort and the other from kilosort2
        # and is updated by the user when doing cluster assignment in phy
        # See comments above this class definition for a bit more info
        if fileExists(self.fname_root, 'cluster_groups.csv'):
            self.cluster_id, self.group = np.loadtxt(
                os.path.join(
                    self.fname_root, 'cluster_groups.csv'),
                unpack=True, skiprows=1, dtype=dtype)
        if fileExists(self.fname_root, 'cluster_group.tsv'):
            self.cluster_id, self.group = np.loadtxt(
                os.path.join(
                    self.fname_root, 'cluster_group.tsv'),
                unpack=True, skiprows=1, dtype=dtype)

        """
        Output some information to the user if self.cluster_id is still None
        it implies that data has not been sorted / curated
        """
        # if self.cluster_id is None:
        #     print(f"Searching {os.path.join(self.fname_root)} and...")
        #     warnings.warn("No cluster_groups.tsv or cluster_group.csv file was found.\
        #         Have you manually curated the data (e.g with phy?")

        # HWPD 20200527
        # load cluster_info file and add X co-ordinate to it
        if fileExists(self.fname_root, 'cluster_info.tsv'):
            self.cluster_info = pd.read_csv(
                os.path.join(
                    self.fname_root, 'cluster_info.tsv'), "\t")
            if fileExists(self.fname_root, 'channel_positions.npy') and \
                    fileExists(self.fname_root, 'channel_map.npy'):
                chXZ = np.load(os.path.join(
                    self.fname_root, 'channel_positions.npy'))
                chMap = np.load(os.path.join(
                    self.fname_root, 'channel_map.npy'))
                chID = np.asarray([np.argmax(
                    chMap == x) for x in self.cluster_info.ch.values])
                self.cluster_info['chanX'] = chXZ[chID, 0]
                self.cluster_info['chanY'] = chXZ[chID, 1]

        dtype = {'names': ('cluster_id', 'KSLabel'), 'formats': ('i4', 'S10')}
        # 'Raw' labels from a kilosort session
        if fileExists(self.fname_root, 'cluster_KSLabel.tsv'):
            self.ks_cluster_id, self.ks_group = np.loadtxt(os.path.join(
                self.fname_root, 'cluster_KSLabel.tsv'), unpack=True,
                skiprows=1, dtype=dtype)
        if fileExists(self.fname_root, 'spike_clusters.npy'):
            self.spk_clusters = np.squeeze(np.load(os.path.join(
                self.fname_root, 'spike_clusters.npy')))
        if fileExists(self.fname_root, 'spike_times.npy'):
            self.spk_times = np.squeeze(np.load(os.path.join(
                self.fname_root, 'spike_times.npy')))
            return True
        warnings.warn("No spike times or clusters were found \
            (spike_times.npy or spike_clusters.npy).\
                You should run KiloSort")
        return False

    def removeNoiseClusters(self):
        """
        Removes clusters with labels 'noise' and 'mua' in self.group
        """
        if self.cluster_id is not None:
            self.good_clusters = []
            for id_group in zip(self.cluster_id, self.group):
                if 'noise' not in id_group[1].decode() and \
                                'mua' not in id_group[1].decode():
                    self.good_clusters.append(id_group[0])

    def removeKSNoiseClusters(self):
        """
        Removes "noise" and "mua" clusters from the kilosort labelled stuff
        """
        for cluster_id, kslabel in zip(self.ks_cluster_id, self.ks_group):
            if 'good' in kslabel.decode():
                self.good_clusters.append(cluster_id)

class OpenEphysBase(FigureMaker):
    """
    Base class for openephys anaylsis with data recorded in either
    the NWB or binary format

    Parameters
    ----------
    pname_root : str
        The top-level directory, typically in form of YYYY-MM-DD_HH-MM-SS

    Notes
    ----
    This isn't really an Abstract Base Class (as with c++) as Python doesn't
    really have this concept but it forms the backbone for two other classes
    (OpenEphysNPX & OpenEphysNWB)
    """
    def __init__(self, pname_root: str, **kwargs):
        super().__init__()
        # top-level directory, typically of form YYYY-MM-DD_HH-MM-SS
        assert(os.path.exists(pname_root))
        self.pname_root = pname_root
        self.settings = None
        self.kilodata = None
        self.rawData = None
        self.xy = None
        self.xyTS = None
        self.recording_start_time = 0
        self.ts = None
        self.ttl_data = None
        self.ttl_timestamps = None
        # a list of np.arrays, nominally containing tetrode data in
        # format nspikes x 4 x 40
        self.spikeData = None
        self.accelerometerData = None
        # This will become an instance of OESettings.Settings
        self.settings = None
        if ('jumpmax' in kwargs):
            self.jumpmax = kwargs['jumpmax']
        else:
            self.jumpmax = 100
        self._ppm = getattr(self, 'ppm', 300)

    @property
    def ppm(self):
        return self._ppm

    @ppm.setter
    def ppm(self, value):
        if hasattr(self, 'PosCalcs'):
            self._ppm = value
            P = getattr(self, 'PosCalcs')
            P.ppm = value
            sample_rate = getattr(P, 'sample_rate')
            P.postprocesspos({'SampleRate': sample_rate})
            setattr(self, 'xy', P.xy)
            setattr(self, 'dir', P.dir)
            setattr(self, 'speed', P.speed)

    def load(self, *args, **kwargs):
        # Overridden by sub-classes
        pass

    def loadKilo(self, **kwargs):
        import os
        if 'pname' in kwargs:
            pname = kwargs['pname']
        else:
            pname = self.pname_root
        # Loads a kilosort session
        kilodata = None
        if pname is not None:
            if os.path.exists(pname):
                # pname_root gets walked through and over-written with
                # correct location of kilosort data
                kilodata = KiloSortSession(pname)
        if kilodata is not None:
            if kilodata.load():
                try:
                    kilodata.removeKSNoiseClusters()
                except Exception:
                    pass
        self.kilodata = kilodata

    def __loadSettings__(self):
        # Loads the settings.xml data
        if self.settings is None:
            # pname_root gets walked through and over-written with
            # correct location of settings.xml
            settings = Settings(self.pname_root)
            settings.parse()
            self.settings = settings

    def __loaddata__(self, **kwargs):
        self.load(self.pname_root, **kwargs)  # some knarly hack

    def __calcTrialLengthFromBinarySize__(
            self, path2file: str, n_channels=384, sample_rate=30000):
        """
        Returns the time taken to run the trial (in seconds) based on the size
        of the binary file on disk
        """
        import os
        status = os.stat(path2file)
        return status.st_size / (2.0 * n_channels * sample_rate)

    def memmapBinaryFile(self, path2file: str, n_channels=384, **kwargs):
        """
        Returns a numpy memmap of the 30Khz sampled high frequency data in the
        file 'continuous.dat', if present
        """
        import os
        if os.path.exists(path2file):
            status = os.stat(path2file)
            n_samples = status.st_size / (2.0 * n_channels)
            mmap = np.memmap(
                path2file, np.int16, 'r', 0, (n_channels, n_samples),
                order='F')
            return mmap

    def exportPos(self):
        xy = self.plotPos(show=False)
        out = np.hstack([xy, self.xyTS[:, np.newaxis]])
        np.savetxt(
            'position.txt', out, delimiter=',',
            fmt=['%3.3i', '%3.3i', '%3.3f'])

    def save_ttl(self, out_fname):
        """
        Saves the ttl data to text file out_fname
        """
        if (len(self.ttl_data) > 0) and (len(self.ttl_timestamps) > 0):
            data = np.array([self.ttl_data, self.ttl_timestamps])
            if data.shape[0] == 2:
                data = data.T
            np.savetxt(out_fname, data, delimiter='\t')

    def filterPosition(self, filter_dict: dict):
        if hasattr(self, 'PosCalcs'):
            P = getattr(self, 'PosCalcs')
            P.filterPos(filter_dict)
            setattr(self, 'xy', P.xy)
            setattr(self, 'dir', P.dir)
            setattr(self, 'speed', P.speed)

    def getClusterSpikeTimes(self, cluster: int):
        '''
        Returns the spike times in seconds of the given cluster
        '''
        spk_times = self.kilodata.spk_times.T
        return spk_times[self.kilodata.spk_clusters == cluster].astype(np.int64)

    def plotSummary(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)
        fig = self.makeSummaryPlot(ts, **kwargs)
        plt.show()
        return fig

    def plotSpikesOnPath(self, cluster: int = None, **kwargs):
        ts = None
        if cluster is not None:
            ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeSpikePathPlot(ts, **kwargs)
        plt.show()
        return ax

    def plotRateMap(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeRateMap(ts)
        plt.show()
        return ax

    def plotHDMap(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeHDPlot(ts, **kwargs)
        plt.show()
        return ax

    def plotSAC(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeSAC(ts, **kwargs)
        plt.show()
        return ax

    def plotSpeedVsRate(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeSpeedVsRatePlot(ts, **kwargs)
        plt.show()
        return ax

    def plotSpeedVsHeadDirection(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)  # in samples
        ax = self.makeSpeedVsHeadDirectionPlot(ts, **kwargs)
        plt.show()
        return ax

    def plotEEGPower(self, channel=0):
        """
        Plots LFP power

        Parameters
        ----------
        channel : int
            The channel from which to plot the power

        See Also
        -----
        ephysiopy.common.ephys_generic.EEGCalcsGeneric.plotPowerSpectrum()
        """
        from ephysiopy.common.ephys_generic import EEGCalcsGeneric
        if self.rawData is None:
            print("Loading raw data...")
            self.load(loadraw=True)
        from scipy import signal
        n_samples = np.shape(self.rawData[:, channel])[0]
        s = signal.resample(self.rawData[:, channel], int(n_samples/3e4) * 500)
        E = EEGCalcsGeneric(s, 500)
        power_res = E.calcEEGPowerSpectrum()
        ax = self.makePowerSpectrum(
            power_res[0], power_res[1], power_res[2],
            power_res[3], power_res[4],
        )
        plt.show()
        return ax

    def plotXCorr(self, cluster: int, **kwargs):
        ts = self.getClusterSpikeTimes(cluster)
        ax = self.makeXCorr(ts)
        plt.show()
        return ax

    def plotPSTH(self, **kwargs):
        """Plots the peri-stimulus time histogram for all the 'good' clusters

        Given some data has been recorded in the ttl channel, this method plots
        the PSTH for each 'good' cluster and just keeps spitting out figure
        windows
        """
        self.__loadSettings__()
        self.settings.parseStimControl()
        if self.kilodata is None:
            self.loadKilo(**kwargs)
        from ephysiopy.common.spikecalcs import SpikeCalcsGeneric
        # in seconds
        spk_times = (self.kilodata.spk_times.T[0] / 3e4) + self.ts[0]
        S = SpikeCalcsGeneric(spk_times)
        # this is because some of the trials have two weird events
        # logged at about 2-3 minutes in...
        S.event_ts = self.ttl_timestamps[2::2]
        S.spk_clusters = self.kilodata.spk_clusters
        S.stim_width = 0.01  # in seconds
        for x in self.kilodata.good_clusters:
            print(next(S.plotPSTH(x)))

    def plotEventEEG(self):
        from ephysiopy.common.ephys_generic import EEGCalcsGeneric
        if self.rawData is None:
            print("Loading raw data...")
            self.load(loadraw=True)
        E = EEGCalcsGeneric(self.rawData[:, 0], 3e4)
        # this is because some of the trials have two weird events
        # logged at about 2-3 minutes in...
        event_ts = self.ttl_timestamps[2::2]
        E.plotEventEEG(event_ts)

class OpenEphysNPX(OpenEphysBase):
    """
    The main class for dealing with data recorded using Neuropixels probes
    under openephys.
    """
    def __init__(self, pname_root: str):
        super().__init__(pname_root)
        self.path2PosData = None
        self.path2APdata = None
        self.path2LFPdata = None
        self.path2syncmessages = None
        self.path2APOEBin = None
        self.path2PosOEBin = None

    def load(
            self, pname_root=None, experiment_name='experiment1',
            recording_name='recording1', **kwargs):
        """
        Loads data recorded in the OE 'flat' binary format.

        Parameters
        ----------
        pname_root : str
            The top level directory, typically in form of YYYY-MM-DD_HH-MM-SS

        recording_name : str
            The directory immediately beneath pname_root

        See Also
        --------
        See open-ephys wiki pages
        """
        self.isBinary = True
        import os
        import re
        APdata_match = re.compile('Neuropix-PXI-[0-9][0-9][0-9].0')
        LFPdata_match = re.compile('Neuropix-PXI-[0-9][0-9][0-9].1')
        PosTracker_match = re.compile('Pos_Tracker-[0-9][0-9][0-9].[0-9]/BINARY_group_[0-9]')
        sync_message_file = None
        self.recording_start_time = None
        ap_sample_rate = getattr(self, 'ap_sample_rate', 30000)

        if pname_root is None:
            pname_root = self.pname_root

        for d, c, f in os.walk(pname_root):
            for ff in f:
                if '.' not in c:  # ignore hidden directories
                    if 'data_array.npy' in ff:
                        if PosTracker_match.search(d):
                            if self.path2PosData is None:
                                self.path2PosData = os.path.join(d)
                            self.path2PosOEBin = Path(d).parents[1]
                    if 'continuous.dat' in ff:
                        if APdata_match.search(d):
                            self.path2APdata = os.path.join(d)
                            self.path2APOEBin = Path(d).parents[1]
                        if LFPdata_match.search(d):
                            self.path2LFPdata = os.path.join(d)
                    if 'sync_messages.txt' in ff:
                        sync_file = os.path.join(
                            d, 'sync_messages.txt')
                        if fileContainsString(sync_file, 'Processor'):
                            sync_message_file = sync_file

        if self.path2PosData is not None:
            pos_data = np.load(os.path.join(
                self.path2PosData, 'data_array.npy'))
            pos_ts = np.load(os.path.join(self.path2PosData, 'timestamps.npy'))
            pos_ts = np.ravel(pos_ts)
            default_pos_sample_rate = np.floor(1/np.mean(np.diff(pos_ts)/ap_sample_rate))
            sample_rate = getattr(self, 'pos_sample_rate', default_pos_sample_rate)
            self.xyTS = pos_ts
            self.pos_sample_rate = sample_rate
            self.orig_x = pos_data[:, 0]
            self.orig_y = pos_data[:, 1]

            P = PosCalcsGeneric(
                pos_data[:, 0], pos_data[:, 1], cm=True, ppm=self.ppm, jumpmax=self.jumpmax)
            xy, hdir = P.postprocesspos({'SampleRate': sample_rate})
            setattr(self, 'PosCalcs', P)
            self.xy = xy
            self.dir = hdir
            self.speed = P.speed
        else:
            warnings.warn("Could not find the pos data. \
                Make sure there is a pos_data folder with data_array.npy and timestamps.npy in")

        n_channels = getattr(self, 'n_channels', 384)
        trial_length = 0  # make sure a trial_length has a value
        if self.path2APdata is not None:
            if fileExists(self.path2APdata, 'continuous.dat'):
                trial_length = self.__calcTrialLengthFromBinarySize__(
                    os.path.join(self.path2APdata, 'continuous.dat'),
                    n_channels, ap_sample_rate)
        # Load the start time from the sync_messages file
        if sync_message_file is not None:
            with open(sync_message_file, 'r') as f:
                sync_strs = f.read()
            sync_lines = sync_strs.split('\n')
            for line in sync_lines:
                if 'subProcessor: 0' in line:
                    idx = line.find('start time: ')
                    start_val = line[idx + len('start time: '):-1]
                    tmp = start_val.split('@')
                    recording_start_time = float(tmp[0]) / float(tmp[1][0:-1])
        else:
            recording_start_time = self.xyTS[0]
        self.recording_start_time = recording_start_time
        # this way of creating timestamps will be fine for single probes
        # but will need to be modified if using multiple probes and/ or
        # different timestamp syncing method
        # OE's 'synchronised_timestamps.npy' should now take care of this
        self.ts = np.arange(
            recording_start_time, trial_length+recording_start_time,
            1.0 / ap_sample_rate)

    def plotSpectrogramByDepth(
            self, nchannels=384, nseconds=100, maxFreq=125,
            **kwargs):
        """
        Plots a heat map spectrogram of the LFP for each channel.

        Line plots of power per frequency band and power on a subset
        of channels are also displayed to the right and above the main plot.

        Parameters
        ----------
        nchannels : int
            The number of channels on the probe
        nseconds : int, optional
            How long in seconds from the start of the trial to do the
            spectrogram for (for speed).
            Default 100
        maxFreq : int
            The maximum frequency in Hz to plot the spectrogram out to.
            Maximum 1250.
            Default 125
        kwargs: legal values are: 'frequencies' and 'channels'; both are lists
                that denote which frequencies to show the mean power of
                along the length of the probe and,
                which channels to show the frequency spectra of.

        Notes
        -----
        Should also allow kwargs to specify exactly which channels and / or
        frequency bands to do the line plots for
        """
        import os
        lfp_file = os.path.join(self.path2LFPdata, 'continuous.dat')
        status = os.stat(lfp_file)
        nsamples = int(status.st_size / 2 / nchannels)
        mmap = np.memmap(
            lfp_file, np.int16, 'r', 0, (nchannels, nsamples), order='F')
        # Load the channel map
        #  Assumes this is in the AP data location and that kilosort was run
        # channel_map = np.squeeze(np.load(os.path.join(
        #     self.path2APdata, 'channel_map.npy')))
        channel_map = np.arange(nchannels)
        lfp_sample_rate = 2500
        data = np.array(mmap[channel_map, 0:nseconds*lfp_sample_rate])
        from ephysiopy.common.ephys_generic import EEGCalcsGeneric
        E = EEGCalcsGeneric(data[0, :], lfp_sample_rate)
        E.calcEEGPowerSpectrum()
        # Select a subset of the full amount of data to display later
        spec_data = np.zeros(shape=(data.shape[0], len(E.sm_power[0::50])))
        for chan in range(data.shape[0]):
            E = EEGCalcsGeneric(data[chan, :], lfp_sample_rate)
            E.calcEEGPowerSpectrum()
            spec_data[chan, :] = E.sm_power[0::50]

        x, y = np.meshgrid(E.freqs[0::50], channel_map)
        import matplotlib.colors as colors
        from matplotlib.pyplot import cm
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        _, spectoAx = plt.subplots()
        if 'cmap' in kwargs:
            cmap = kwargs['cmap']
        else:
            cmap = 'bone'
        spectoAx.pcolormesh(
            x, y, spec_data, edgecolors='face', cmap=cmap,
            norm=colors.LogNorm(), shading='nearest')
        if 'minFreq' in kwargs:
            minFreq = kwargs['minFreq']
        else:
            minFreq = 0
        spectoAx.set_xlim(minFreq, maxFreq)
        spectoAx.set_ylim(channel_map[0], channel_map[-1])
        spectoAx.set_xlabel('Frequency (Hz)')
        spectoAx.set_ylabel('Channel')
        divider = make_axes_locatable(spectoAx)
        channel_spectoAx = divider.append_axes(
            "top", 1.2, pad=0.1, sharex=spectoAx)
        meanfreq_powerAx = divider.append_axes(
            "right", 1.2, pad=0.1, sharey=spectoAx)
        plt.setp(
            channel_spectoAx.get_xticklabels() +
            meanfreq_powerAx.get_yticklabels(), visible=False)

        mn_power = np.mean(spec_data, 0)
        if 'channels' in kwargs:
            channels = kwargs['channels']
            cols = iter(cm.rainbow(
                np.linspace(0, 1, (len(kwargs['channels'])))))
        else:
            channels = np.arange(0, spec_data.shape[0], 60)
            cols = iter(cm.rainbow(
                np.linspace(0, 1, (nchannels//60)+1)))
        for i in channels:
            c = next(cols)
            channel_spectoAx.plot(
                E.freqs[0::50], 10*np.log10(spec_data[i, :]/mn_power),
                c=c, label=str(i))

        channel_spectoAx.set_ylabel('Channel power(dB)')
        channel_spectoAx.legend(
            bbox_to_anchor=(0., 1.02, 1., .102), loc='lower left',
            mode='expand', fontsize='x-small', ncol=4)

        if 'frequencies' in kwargs:
            lower_freqs = kwargs['frequencies']
            upper_freqs = lower_freqs[1::]
            inc = np.diff([lower_freqs[-2], lower_freqs[-1]])
            upper_freqs = np.append(upper_freqs, lower_freqs[-1]+inc)
        else:
            freq_inc = 6
            lower_freqs = np.arange(1, maxFreq-freq_inc, freq_inc)
            upper_freqs = np.arange(1+freq_inc, maxFreq, freq_inc)

        cols = iter(cm.nipy_spectral(np.linspace(0, 1, len(upper_freqs))))
        mn_power = np.mean(spec_data, 1)
        print(f"spec_data shape = {np.shape(spec_data)}")
        print(f"mn_power shape = {np.shape(mn_power)}")
        for freqs in zip(lower_freqs, upper_freqs):
            freq_mask = np.logical_and(
                E.freqs[0::50] > freqs[0],
                E.freqs[0::50] < freqs[1])
            mean_power = 10*np.log10(np.mean(
                spec_data[:, freq_mask], 1)/mn_power)
            c = next(cols)
            meanfreq_powerAx.plot(
                mean_power, channel_map, c=c, label=str(
                    freqs[0]) + " - " + str(freqs[1]))
        meanfreq_powerAx.set_xlabel('Mean freq. band power(dB)')
        meanfreq_powerAx.legend(
            bbox_to_anchor=(0., 1.02, 1., .102), loc='lower left',
            mode='expand', fontsize='x-small', ncol=1)
        if 'saveas' in kwargs:
            saveas = kwargs['saveas']
            plt.savefig(saveas)
        plt.show()

class OpenEphysNWB(OpenEphysBase):
    """
    Parameters
    ------------
    pname_root : str
        The top level directory, typically in form of YYYY-MM-DD_HH-MM-SS
    """

    def __init__(self, pname_root, **kwargs):
        super().__init__(pname_root)
        self.nwbData = None  # handle to the open nwb file (HDF5 file object)
        self.rawData = None  # np.array holding the raw, continuous recording
        self.recording_name = None  # the recording name ('recording0' etc)
        self.isBinary = False
        self.xy = None

    def load(
            self, pname_root: None, session_name=None, recording_name=None,
            loadraw=False, loadspikes=False, savedat=False):
        """
        Loads xy pos from binary part of the hdf5 file and data resulting from
        a Kilosort session (see KiloSortSession class above)

        Parameters
        ----------
        pname_root : str
            The top level directory, typically the one named
            YYYY-MM-DD_HH-MM-SS
            NB In the nwb format this directory contains the experiment_1.nwb
            and settings.xml files
        session_name : str
            Defaults to experiment_1.nwb
        recording_name : str
            Defaults to recording0
        loadraw : bool
            Defaults to False; if True will load and save the
            raw part of the data
        savedat : bool
            Defaults to False; if True will extract the electrode
            data part of the hdf file and save as 'experiment_1.dat'
            NB only works if loadraw is True. Also note that this
            currently saves 64 channels worth of data (ie ignores
            the 6 accelerometer channels)
        """

        import h5py
        import os
        if pname_root is None:
            pname_root = self.pname_root
        if session_name is None:
            session_name = 'experiment_1.nwb'
        self.nwbData = h5py.File(os.path.join(
            pname_root, session_name), mode='r')
        # Position data...
        if self.recording_name is None:
            if recording_name is None:
                recording_name = 'recording1'
            self.recording_name = recording_name
        try:
            self.xy = np.array(
                self.nwbData['acquisition']['timeseries']
                [self.recording_name]['events']['binary1']['data'])

            self.xyTS = np.array(
                self.nwbData['acquisition']['timeseries']
                [self.recording_name]['events']['binary1']['timestamps'])
            self.xy = self.xy[:, 0:2]
        except Exception:
            self.xy = None
            self.xyTS = None
        try:
            # TTL data...
            self.ttl_data = np.array(
                self.nwbData['acquisition']['timeseries']
                [self.recording_name]['events']['ttl1']['data'])
            self.ttl_timestamps = np.array(
                self.nwbData['acquisition']['timeseries']
                [self.recording_name]['events']['ttl1']['timestamps'])
        except Exception:
            self.ttl_data = None
            self.ttl_timestamps = None

        # ...everything else
        try:
            self.__loadSettings__()
            fpgaId = self.settings.processors['Sources/Rhythm FPGA'].NodeId
            fpgaNode = 'processor' + str(fpgaId) + '_' + str(fpgaId)
            self.ts = np.array(
                self.nwbData['acquisition']['timeseries']
                [self.recording_name]['continuous'][fpgaNode]['timestamps'])
            if loadraw is True:
                print("Attempting to reference ALL the raw data...")
                self.rawData = self.nwbData['acquisition']['time\
                    series'][self.recording_name]['\
                        continuous'][fpgaNode]['data']
                print("Referenced the raw data! Yay!\nParsing channels...")
                self.settings.parseProcessor()  # get the neural data channels
                print("Channels parsed\nAccessing accelerometer data...")
                self.accelerometerData = self.rawData[:, 64:]
                print("Accessed the accelerometer data\
                    \nAttempting to access the raw data...")
                self.rawData = self.rawData[:, 0:64]
                print("Got the raw data!")
                if (savedat is True):
                    data2save = self.rawData[:, 0:64]
                    data2save.tofile(os.path.join(
                        pname_root, 'experiment_1.dat'))
            if loadspikes is True:
                if self.nwbData['acquisition']['timeseries'][self.recording_name]['\
                        spikes']:
                    # Create a dictionary containing keys 'electrode1',
                    # 'electrode2' etc and None for values
                    electrode_dict = dict.fromkeys(
                        self.nwbData['acquisition']['\
                            timeseries'][self.recording_name]['spikes'].keys())
                    # Each entry in the electrode dict is a dict
                    #  containing keys 'timestamps' and 'data'...
                    for i_electrode in electrode_dict.keys():
                        data_and_ts_dict = {'timestamps': None, 'data': None}
                        data_and_ts_dict['timestamps'] = np.array(
                            self.nwbData['acquisition']['\
                                timeseries'][self.recording_name]['\
                                    spikes'][i_electrode]['timestamps'])
                        data_and_ts_dict['data'] = np.array(
                            self.nwbData['acquisition']['\
                                timeseries'][self.recording_name]['\
                                    spikes'][i_electrode]['data'])
                        electrode_dict[i_electrode] = data_and_ts_dict
                self.spikeData = electrode_dict
        except Exception:
            self.ts = self.xy

class OpenEphysBinary(OpenEphysBase):
    """
    The main class for dealing with data recorded using openephys
    and the Rhythm-FPGA module .
    """
    def __init__(self, pname_root):
        super().__init__(pname_root)
        self.path2PosData = None
        self.path2APdata = None
        self.path2LFPdata = None
        self.rawData = None

    def load(
            self, pname_root=None, experiment_name='experiment1',
            recording_name='recording1', loadraw=False, n_channels=64):
        """
        Loads data recorded in the OE 'flat' binary format.

        Parameters
        ----------
        pname_root : str
            The top level directory, typically in form of YYYY-MM-DD_HH-MM-SS

        recording_name : str
            The directory immediately beneath pname_root

        See Also
        --------
        See open-ephys wiki
        """
        self.isBinary = True
        import os
        import re
        APdata_match = re.compile('Rhythm_FPGA-[0-9][0-9][0-9].0')
        LFPdata_match = re.compile('Rhythm_FPGA-[0-9][0-9][0-9].1')
        sync_message_file = None
        self.recording_start_time = None
        ap_sample_rate = getattr(self, 'ap_sample_rate', 30000)

        if pname_root is None:
            pname_root = self.pname_root

        for d, c, f in os.walk(pname_root):
            for ff in f:
                if '.' not in c:  # ignore hidden directories
                    if 'data_array.npy' in ff:
                        self.path2PosData = os.path.join(d)
                    if 'continuous.dat' in ff:
                        if APdata_match.search(d):
                            self.path2APdata = os.path.join(d)
                        if LFPdata_match.search(d):
                            self.path2LFPdata = os.path.join(d)
                    if 'sync_messages.txt' in ff:
                        sync_file = os.path.join(
                            d, 'sync_messages.txt')
                        if fileContainsString(sync_file, 'Processor'):
                            sync_message_file = sync_file

        if self.path2PosData is not None:
            pos_data = np.load(os.path.join(
                self.path2PosData, 'data_array.npy'))
            pos_ts = np.load(os.path.join(
                self.path2PosData, 'timestamps.npy'))
            pos_ts = np.ravel(pos_ts)
            sample_rate = np.floor(1/np.mean(np.diff(pos_ts)/ap_sample_rate))
            self.xyTS = pos_ts
            self.pos_sample_rate = sample_rate
            self.orig_x = pos_data[:, 0]
            self.orig_y = pos_data[:, 1]

            P = PosCalcsGeneric(
                pos_data[:, 0], pos_data[:, 1], cm=True, ppm=self.ppm, jumpmax=self.jumpmax)
            xy, hdir = P.postprocesspos({'SampleRate': sample_rate})
            setattr(self, 'PosCalcs', P)
            self.xy = xy
            self.dir = hdir
            self.speed = P.speed

        n_channels = getattr(self, 'n_channels', 384)
        trial_length = 0  # make sure a trial_length has a value
        if fileExists(self.path2APdata, 'continuous.dat'):
            trial_length = self.__calcTrialLengthFromBinarySize__(
                os.path.join(self.path2APdata, 'continuous.dat'),
                n_channels, ap_sample_rate)

        if loadraw is True:
            if fileExists(self.path2APdata, 'continuous.dat'):
                status = os.stat(os.path.join(
                    self.path2APdata, 'continuous.dat'))
                n_samples = int(status.st_size / 2 / n_channels)
                mmap = np.memmap(os.path.join(
                    self.path2APdata, 'continuous.dat'),
                    np.int16, 'r', 0, (n_channels, n_samples), 'C')
                self.rawData = np.array(mmap, dtype=np.float64)

        # Load the start time from the sync_messages file
        recording_start_time = 0
        if sync_message_file is not None:
            with open(sync_message_file, 'r') as f:
                sync_strs = f.read()
            sync_lines = sync_strs.split('\n')
            for line in sync_lines:
                if 'subProcessor: 0' in line:
                    idx = line.find('start time: ')
                    start_val = line[idx + len('start time: '):-1]
                    tmp = start_val.split('@')
                    recording_start_time = float(tmp[0]) / float(tmp[1][0:-1])
        else:
            recording_start_time = self.xyTS[0]
        self.recording_start_time = recording_start_time
        self.ts = np.arange(
            recording_start_time, trial_length+recording_start_time,
            1.0 / ap_sample_rate)


from collections import defaultdict
from itertools import product
import pandas as pd
import numpy as np

def filt(d, args):

    d2 = defaultdict(lambda: defaultdict(list))
    for pos, calls in d.items():
        for k, v in calls.items():
            if type(v) != list:
                d2[pos][k] = v
            else:
                for i, option in enumerate(v):
                    p = args.P_value
                    alt = calls.get('tumour_alts')[i]
                    percent_norm = calls.get('read_count_normal')[i][0] /\
                        calls.get('read_count_normal')[i][1]
                    if percent_norm > 0.07:
                        continue
                    if alt[0] in ['-', '+']:
                        if percent_norm > 0.00:
                            continue
                    if alt == calls.get('ref'):  # any LOH that slip through that arent interesting
                        continue
                    if calls.get('read_count_tumour')[i][1] > 120 and \
                            calls.get('read_count_normal')[i][1] > 60:
                        p *= 0.001  # if in CNV requires lower P
                    if calls.get('read_count_tumour')[i][1] > 480 and \
                            calls.get('read_count_normal')[i][1] > 240:
                        p *= 0.00001  # if in CNV requires lower P
                    if 0.07 > calls.get('strand')[i] or \
                            calls.get('strand')[i] > 0.93:
                        continue
                    if float(calls.get('tumour_P')[i]) > p:
                        continue
                    if calls.get('read_count_tumour')[i][1] < args.min_depth_tumour or \
                        calls.get('read_count_normal')[i][1] < args.min_depth_normal:
                        break  # HOM dels that actually line up properly should have read depth zero
                    if all(map(lambda x: x > args.min_mean_base_quality, [
                            float(calls.get('tumour_qual')[i]),
                            calls.get('tumour_qual_all'),
                            calls.get('normal_qual_all')])):
                        d2[pos][k].append(option)
        if not d2[pos].get('tumour_alts'):
            del d2[pos]
    df = pd.DataFrame.from_dict(d2, orient='index')

    if not len(df):
        return

    # make pseudo-VCF
    # TODO make a proper VCF
    df['tmp_index_col'] = df.index
    df[['chrom', 'pos']] = df.tmp_index_col.str.split(':', expand=True)
    df.drop('tmp_index_col', axis=1, inplace=True)
    cols = df.columns.tolist()
    cols = cols[-2:] + cols[:-2]
    df = df[cols]

    if args.MNVs:
        df = MNV(df)

    # multiple alts as separate records
    df = df.explode(['tumour_alts',
                     'tumour_P', 'tumour_qual', 'normal_qual', 'strand',
                     'read_count_tumour', 'read_count_normal'])

    fout_name = '_'.join([
        args.output,
        str(args.P_value).split('.')[1],
        str(args.min_mean_base_quality),
        str(args.min_depth_tumour),
        str(args.min_depth_normal)])

    df.to_csv(fout_name + '.tsv', index=False, sep='\t')

def distance(r1, r2):
    """
    Return r2['pos'] - r1['pos'] if r1['chrom'] == r2['chrom']
    (corollary, return 0 iff r1 and r2 have same position)
    Return np.inf if r1['chrom'] != r2['chrom']
    """
    if r1.get('chrom') is None or r1.get('chrom') != r2.get('chrom'):
        return np.inf
    return r2['pos'] - r1['pos']

def mergepos(altalleles):
    """
    Given a list-of-lists representation of (possibly poly-allelic)
    contiguous SNVs, return a list-of-sequences representation of the MNV
    alleles at the first position

    Args:
        altalleles: list of lists, the inner lists representing allele(s)
            at first and subsequent contiguous positions of a variant.
            For example:
            [['A']]                         # mono-allelic SNV
            [['A', 'C']]                    # poly-allelic SNV
            [['A'], ['C']]                  # mono-allelic MNV
            [['A', 'C'], ['A'], ['C', 'G']] # poly-allelic MNV
    Returns:
        list of MNV alleles at the first position
        >>> mergepos([['A']])
        ['A']
        >>> mergepos([['A']])
        ['A', 'C']
        >>> mergepos([['A'], ['C']])
        ['AC']
        >>> # note: the result is Cartesian product
        >>> mergepos([['A', 'C'], ['A'], ['C', 'G']])
        ['AAC', 'AAG', 'CAC', 'CAG']
    """
    if not altalleles:
        return []
    return [''.join(bases) for bases in product(*altalleles)]

def mergesnvs(snvs):
    """
    Takes a list of contiguous SNV records with a list-of-alleles
    representation in tumour_alts and returns MNV starting at the first
    position.

    Returns:
        Empty list if snvs is empty, otherwise a single-element list
        containing the MNV record

    Raises:
        ValueError if SNV positions are not contiguous
    """
    if not snvs:
        return []

    posx = [snv['pos'] for snv in snvs]
    posx0 = [pos - posx[0] for pos in posx]
    if not posx0 == list(range(len(snvs))):
        raise ValueError(f'{posx} are not contiguous')

    tumour_alts = mergepos([snv['tumour_alts'] for snv in snvs])
    n_alleles = len(tumour_alts)
    # simplification: in combining contiguous SNV calls into MNVs it's not
    # immediately obvious what to do with the different values for tumour_P,
    # tumour_qual, normal_qual, strand, read_count_tumour, read_count_normal.
    # Currently we take the simplest possible approach by repeating the first
    # value from the first position only. Other simple approaches include
    # taking average values or worst values across positions and alleles.
    first = snvs[0]
    return [{
        **snvs[0],
        'tumour_alts':       tumour_alts,
        'ref':               ''.join(snv['ref'] for snv in snvs),
        'normal_alts':       ''.join(snv['normal_alts'] for snv in snvs),
        'tumour_P':          [first['tumour_P'][0]] * n_alleles,
        'tumour_qual':       [first['tumour_qual'][0]] * n_alleles,
        'normal_qual':       [first['normal_qual'][0]] * n_alleles,
        'strand':            [first['strand'][0]] * n_alleles,
        'read_count_tumour': [first['read_count_tumour'][0]] * n_alleles,
        'read_count_normal': [first['read_count_normal'][0]] * n_alleles
    }]

def MNV(df):
    """
    Convert any contiguous SNVs to MNVs.
    """
    df.pos = df.pos.astype(int)
    df.sort_values(by=['chrom', 'pos'])
    records = df.to_dict('records')

    merged = []
    prev = {}
    contiguous = []
    for record in records:
        d = distance(prev, record)
        if d == 1:
            contiguous.append(record)
        else:
            assert d > 1, 'records must be sorted by (chrom, pos) ascending'
            merged.extend(mergesnvs(contiguous))
            contiguous = [record]
        prev = record
    merged.extend(mergesnvs(contiguous))

    return df.from_records(
        merged, index=[f'{x["chrom"]}:{x["pos"]}' for x in merged])


import os
from pathlib import Path
from functools import partial

from langdetect import detect
import numpy as np
from sklearn.preprocessing import normalize
from summa.pagerank_weighted import pagerank_weighted_scipy as _pagerank
from summa.commons import remove_unreachable_nodes as _remove_unreachable_nodes
from summa.summarizer import _set_graph_edge_weights, _add_scores_to_sentences
from summa.commons import build_graph as _build_graph

from text_cleaning_en import clean_text_by_sentences as _clean_text_by_sentences
from summa_score_sentences_use import cut_sentences_by_rule, cosine_similarity

MODEL_PATH = Path(os.environ["LASER"]) / "models/"

def attach_sentence_embeddings(lang, sentences, batch_size=32):
    from laser.shortcuts import lines_to_embeddings

    # don't use extremely short sentences
    sentences_subset = [x for x in sentences if len(x.text) > 5]
    sentence_embeddings_subset = lines_to_embeddings(
        lang,
        [x.text for x in sentences_subset],
        str(MODEL_PATH / "bilstm.93langs.2018-12-26.pt"),
        str(MODEL_PATH / "93langs.fcodes"),
        use_cpu=False,
        batch_size=batch_size
    ).reshape(len(sentences_subset), -1)
    # print(len(sentences_subset))
    # print(sentence_embeddings_subset.shape)
    # l2 normalize
    sentence_embeddings_subset = normalize(
        sentence_embeddings_subset, norm="l2", axis=1, copy=False, return_norm=False
    )
    sentence_embeddings = np.zeros(
        (len(sentences), sentence_embeddings_subset.shape[1]), dtype="float32")
    # A rather hacky way to attach embeddings
    for i, sentence in enumerate(sentences_subset):
        sentence_embeddings[sentence.token, :] = (
            sentence_embeddings_subset[i, :])
    similarities = sentence_embeddings @ sentence_embeddings.T
    # print(similarities[np.tril_indices(similarities.shape[0], k=-1)])
    # print(np.where(np.tril(similarities, k=-1) > 0.95))
    return similarities

def summarize(text, additional_stopwords=None):
    if not isinstance(text, str):
        raise ValueError("Text parameter must be a Unicode object (str)!")

    lang = detect(text)[:2]
    if lang == "en":
        paragraphs = text.split("\n")
        sentences = []
        paragraph_index = 0
        for paragraph in paragraphs:
            # Gets a list of processed sentences.
            if paragraph:
                tmp = _clean_text_by_sentences(
                    paragraph, additional_stopwords)
                if tmp:
                    for j, sent in enumerate(tmp):
                        sent.paragraph = paragraph_index
                        # Hacky way to overwrite token
                        sent.token = len(sentences) + j
                    sentences += tmp
                    paragraph_index += 1
    elif lang == "zh" or lang == "ko":  # zh-Hant sometimes got misclassified into ko
        sentences = cut_sentences_by_rule(text)
    elif lang == "ja":
        raise NotImplementedError("No ja support yet.")
    else:
        return ["Language not suppored! (supported languages: en, zh)"], None, lang
    similarities = attach_sentence_embeddings(
        lang, sentences,  batch_size=32)
    graph = _build_graph([x.token for x in sentences])
    _set_graph_edge_weights(graph, partial(cosine_similarity, similarities))

    # Remove all nodes with all edges weights equal to zero.
    _remove_unreachable_nodes(graph)

    # PageRank cannot be run in an empty graph.
    if len(graph.nodes()) == 0:
        return []

    # Ranks the tokens using the PageRank algorithm. Returns dict of sentence -> score
    pagerank_scores = _pagerank(graph)

    # Adds the summa scores to the sentence objects.
    _add_scores_to_sentences(sentences, pagerank_scores)

    # Sorts the sentences
    sentences.sort(key=lambda s: s.score, reverse=True)
    return sentences, graph, lang

if __name__ == "__main__":
    res, _, lang = summarize("""
By Wednesday morning, he was trying to marry those two thoughts into a single message — both embracing the report and trashing it. “The Mueller Report, despite being written by Angry Democrats and Trump Haters, and with unlimited money behind it ($35,000,000), didn’t lay a glove on me,” he wrote. “I DID NOTHING WRONG.”

In subsequent tweets, he tried again to claim victory amid his victimhood, casting the investigation as a contest in which he prevailed. In terms rarely used regarding a criminal investigation, he asserted that “We waited for Mueller and WON” and denounced “the Witch Hunt, which I have already won.”""")
    assert lang == "en"
    for row in res:
        print(f"{row.score} {row.text}")


from PIL import Image
from keras.models import model_from_json
import os
import utils
import numpy as np

# set filepath
datasets=["DRIVE","STARE"]
fundus_dir="../data/{}/test/images/"
mask_dir="../data/{}/test/mask/"
out_dir="../inference_outputs/{}"
f_model="../pretrained/{}_best.json"
f_weights="../pretrained/{}_best.h5"

for dataset in datasets:
    # make directory
    if not os.path.isdir(out_dir.format(dataset)):
        os.makedirs(out_dir.format(dataset))

    # load the model and weights
    with open(f_model.format(dataset), 'r') as f:
        model=model_from_json(f.read())
    model.load_weights(f_weights.format(dataset))

    # iterate all images
    img_size=(640,640) if dataset=="DRIVE" else (720,720)
    ori_shape=(1,584,565) if dataset=="DRIVE" else (1,605,700)  # batchsize=1
    fundus_files=utils.all_files_under(fundus_dir.format(dataset))
    mask_files=utils.all_files_under(mask_dir.format(dataset))
    for index,fundus_file in enumerate(fundus_files):
        print "processing {}...".format(fundus_file)
        # load imgs
        img=utils.imagefiles2arrs([fundus_file])
        mask=utils.imagefiles2arrs([mask_files[index]])

        # z score with mean, std (batchsize=1)
        mean=np.mean(img[0,...][mask[0,...] == 255.0],axis=0)
        std=np.std(img[0,...][mask[0,...] == 255.0],axis=0)
        img[0,...]=(img[0,...]-mean)/std

        # run inference
        padded_img=utils.pad_imgs(img, img_size)
        vessel_img=model.predict(padded_img,batch_size=1)*255
        cropped_vessel=utils.crop_to_original(vessel_img[...,0], ori_shape)
        final_result=utils.remain_in_mask(cropped_vessel[0,...], mask[0,...])
        Image.fromarray(final_result.astype(np.uint8)).save(os.path.join(out_dir.format(dataset),os.path.basename(fundus_file)))