# -*- coding: utf-8 -*-
"""
Module for the data classes:
    Maps are made of X,Y and associated DataMap(s), while
    Profiles are made of R and associated DataProfile(s)
"""

# External modules
import numpy as np
from numpy import deg2rad

# Units
import astropy.units as u

# various modules from local lib
from .maps_grammar import (analyse_maps_kwargs, analyse_data_kwargs,
                           default_data_separator, list_Data_attr,
                           get_dmap_info_from_dtype, _add_density_suffix,
                           _is_flag_density, remap_suffix, _add_density_prefix)

from .misc_io import (add_suffix, add_prefix, default_float, add_err_prefix,
                      AttrDict, remove_suffix, guess_stepxy, get_extent,
                      cover_linspace)

from . import check, transform

default_data_names = ["data", "edata"]
dict_units = {"XY": u.arcsec, "R": u.arcsec}

class DataMap(object):
    """Data class representing a specific map

    Attributes
    ----------
    data
    edata
    order
    dname
    dunit
    flag
    """
    def __init__(self, **kwargs):
        """
        Args:
            data: numpy array
                Input datas values.
            edata: numpy array [None]
                Uncertainties for the data.
            order: int [0]
                order of the velocity moment. Can be -1 for 'others' (grid)
            dname: str [None]
                Name of the datamap
            flag: str [None]
            dtype: str [""]
            dunit: astropy unit [None]
        """
        # Using the key list in list_Data_attr to set the attributes
        for keyword in list_Data_attr:
            setattr(self, keyword, kwargs.pop(keyword, None))

    def add_transformed_data(self, data, edata=None, suffix=""):
        """Add a transformed map - e.g, deprojected - using a suffix

        Args:
            data:
            edata:
            suffix:

        Returns:

        """
        if suffix == "" or suffix is None:
            print("[add_transformed_data] Failed to add data, suffix is empty")
            return
        new_data_attr = add_suffix("data", suffix)
        setattr(self, new_data_attr, data)
        new_edata_attr = add_err_prefix(new_data_attr)
        setattr(self, new_edata_attr, edata)

    def _reshape_datamap(self, shape):
        """reshape the data
        Args:
            shape:

        """
        self.data = self.data.reshape(shape)
        if self.edata is not None:
            self.edata = self.edata.reshape(shape)

    def deproject_velocities(self, inclin=90.0):
        """Deproject Velocity map and add it
        """

        if self.order != 1:
            print("ERROR: data are not of order 1 [velocities] -- Aborting")
            return
        Vdep, eVdep = transform.deproject_velocities(self.data,
                                                     self.edata,
                                                     inclin)
        self.add_transformed_data(Vdep, eVdep, "dep")

class DataProfile(DataMap):
    def __init__(self, **kwargs):
        """Create a profile class with some data and errors.

        Args:
            data:
            edata:
            **kwargs: (see DataMap)
                order: int
                mname: str
                flag: str
                dtype: ""
        """
        # Using DataMap class attributes
        super().__init__(**kwargs)

        # Now 1D in case the input is 2D
        if self.data is not None:
            self.data = self.data.ravel()
        if self.edata is not None:
            self.edata = self.edata.ravel()

class Map(object):
    """A Map is a set of DataMaps associated with a location grid (X, Y)
    It is used to describe a set of e.g., velocity fields, flux maps, etc
    A grid is associated natively to these DataMaps as well as an orientation
    on the sky.
    If no grid is provided, the grid is set to integer numbers (pixels).

    Attributes
    ----------
    NE_direct: bool [True]
        If True, direct sense for NE meaning East is counter-clockwise
        from North in the map.
    alpha_North: float [0.]
        Angle of the North w.r.t. top. Positive means counter-clockwise.
    X, Y: numpy arrays [None, None]
        Input location grid
    data: numpy array [None]
        Input values
    edata: numpy array [None]
        Uncertainties
    order: int [0]
        Order of the velocity moment. -1 for 'others' (e.g., X, Y)
    """
    def __init__(self, X=None, Y=None, mname=None, mtype="", comment="", **kwargs):
        """
        Args:
            X: numpy array [None]
                Input X axis location array
            Y: numpy array [None]
                Input Y axis location array
            Xcen, Ycen: float, float
                Centre for the X and Y axes. Default is the centre of the image.
            mname: str [None]
                Name of the dataset
            comment: str [""]
                Comment attached to the Map

            **kwargs:
                Any of these attributes can be provided with
                a suffix. E.g., "dataCO" will be understood
                as data and with a flag=CO.
                data: array
                edata: array [None]
                    Uncertainty map
                mname: str
                    Name of the map
                mtype: str
                    Type of the map
                flag: str
                    Flag for the map
                order: int
                    Order for the datamap
        """
        # Empty dictionary for the moments
        self.dmaps = AttrDict()

        # ----------- First the default attributes ----------------
        # We assume that by default the units are the default arcsec
        self.XYunit = kwargs.pop("XYunit", dict_units['XY'])

        # Filling value
        self._fill_value = kwargs.pop("fill_value", 'nan')
        # Method
        self._method = kwargs.pop("method", "linear")

        # Comment, mname and map type
        self.comment = comment
        self.mname = mname
        self.mtype = mtype
        self.overwrite = kwargs.pop("overwrite", False)
        self._force_dtype = kwargs.pop("force_dtype", False)

        # ----------------- End of default attributes -------------

        # ----------- Analyse the kwargs ----------------------
        dict_dmaps, map_kwargs = analyse_data_kwargs(**kwargs)

        # ------------ X and Y for the map ---------------
        # First getting the shape of the data
        # Using X first
        if X is not None:
            self.shape = X.shape
        # If no X, try with the data
        else:
            # Look for the first data which exists
            data = None
            for name_dmap in dict_dmaps.keys():
                if "data" in dict_dmaps[name_dmap]:
                    data = dict_dmaps[name_dmap]['data']
                    break

            if check._check_ifarrays([data]):
                    self.shape = data.shape
            else:
                raise ValueError("No reference shape is provided "
                      "for map '{}'(via X, data) - Ignoring this map".format(
                       mname))
                return

        # Initialise the X, Y coordinates
        self.Xcen = kwargs.pop("Xcen", 0.)
        self.Ycen = kwargs.pop("Ycen", 0.)
        # Pixel scale for the X, Y coordinates
        self.pixel_scale = kwargs.pop("pixel_scale", 1.)
        if not self._init_XY(X, Y):
            raise ValueError("ERROR: X and Y are not compatible")
            return

        # Boolean to say whether E is direct from N (counter-clockwise)
        self.NE_direct = kwargs.pop("NE_direct", True)
        # Angle (degrees) between North and top
        self.alpha_North = kwargs.pop("alpha_north", 0.)
        # Get the matrix for the alignment with North-up
        self.align_xy_NorthEast()

        # --- Attaching each data map in turn -----------
        # Add each datamap one by one
        for dmap_name, dmap_kwargs in dict_dmaps.items():
            if 'dname' not in dmap_kwargs:
                dmap_kwargs['dname'] = dmap_name

            # Test if we need to force dtype
            self.add_datamap(**dmap_kwargs)

    def __getattr__(self, mname):
        for suffix in default_data_names:
            if mname.startswith(suffix):
                for mapname in self.dmaps.keys():
                    if mapname in mname:
                        basename = remove_suffix(mname, mapname,
                                                 separator=default_data_separator)
                        return getattr(self.dmaps[mapname], basename)
        raise AttributeError("'Map' object has no attribute {}".format(mname))

    def __dir__(self):
        return  super().__dir__() + [add_suffix(attr, map) for item in ['data', 'edata']
                for map in self.dmaps.keys() for attr in self.dmaps[map].__dir__()
                if attr.startswith(item)]

    @property
    def ndatamaps(self):
        return len(self.dmaps)

    def _init_XY(self, X, Y):
        """Initialise X and Y

        Args:
            X: numpy array
            Y: numpy array
                Input X, Y grid.
        """
        # Define the grid in case X, Y not yet defined
        # If it is the case, using the reference Map
        if X is None or Y is None:
            # We get the grid in pixel
            print("WARNING: X or Y not provided. Using Pixel XY grid.")
            ref_ind = np.indices(self.shape, dtype=default_float)
            self.X = ref_ind[1] - self.Xcen
            self.Y = ref_ind[0] - self.Ycen
            # And now convert to default unit
            self._convert_to_xyunit()
        else:
            # if X, Y pre-defined, unit is pre-defined too
            if not check._check_consistency_sizes([X, Y]):
                print("ERROR: errors on sizes of X and Y")
                return False
            # Just removing the centre to get 0,0
            self.X = X - self.Xcen
            self.Y = Y - self.Ycen

        # Making sure the shapes agree
        self.X = self.X.reshape(self.shape)
        self.Y = self.Y.reshape(self.shape)
        return True

    def add_datamap(self, **kwargs):
        """Add a new DataMap to the present Map. Will check if
        grid is compatible.

        Args:
            data: 2d array
            order: int
            edata: 2d array
            dname: str
            dtype: str
            flag: str
            dunit: astropy unit
        """
        # Input dname to define the data. If none, define using the counter
        dname = kwargs.pop("dname", None)
        if dname is None or dname=="":
            dname = "{0}{1:02d}".format(self.mname, self.ndatamaps+1)

        data = kwargs.pop("data", None)
        if data is None:
            print("WARNING[attach_data/Map]: cannot attach data: "
                  "it is 'None' (dname is {}) - Ignoring".format(dname))
            return

        if not check._check_ifarrays([data]):
            print("WARNING[attach_data/Map]: these date are not an array"
                  "(dname is {}) - Ignoring".format(dname))
            return

        overwrite = kwargs.pop("overwrite", self.overwrite)

        if self._has_datamap(dname) and not overwrite:
            print("WARNING[attach_data]: data map {} already exists "
                  "- Aborting".format(dname))
            print("WARNING[attach_data]: use overwrite option to force.")
            return

        # Check if we wish to force the dtype / dunit
        force_dtype = kwargs.pop("force_dtype", self._force_dtype)
        if force_dtype:
            dtype = kwargs.get("dtype", None)
            dmap_info = get_dmap_info_from_dtype(dtype, dname)
            # Transfer
            for key, value in dmap_info.items():
                kwargs[key] = value

        self.attach_datamap(DataMap(data=data, dname=dname, **kwargs))

    @property
    def eq_pscale(self):
        return u.pixel_scale(self.pixel_scale * self.XYunit / u.pixel)

    @property
    def _get_pixel_scale(self):
        return (1. * u.pixel).to(self.XYunit, equivalencies=self.eq_pscale).value

    def _convert_to_xyunit(self):
        """Convert XYunit into the default one
        a priori arcseconds.
        """
        self.X *= self.xyunit_per_pixel * self.pixel_scale
        self.Y *= self.xyunit_per_pixel * self.pixel_scale
        # Update the unit
        self.pixel_scale = 1.0
        self.XYunit = dict_units['XY']

    @property
    def xyunit_per_pixel(self):
        return (1. * self.XYunit).to(dict_units['XY'],
                                 self.eq_pscale).value

    def _get_datamap(self, dname=None, order=None):
        """Get the datamap if it exists, and
        check the order

        Args:
            dname:
            order:

        Returns:

        """
        if dname is None:
            if order is None:
                # then just get the first map
                dname = list(self.dmaps.keys())[0]
            else:
                # Then get the first map of right order
                for key in self.dmaps.keys():
                    if self.dmaps[key].order == order:
                        dname = key
                        break

        if self._has_datamap(dname):
            return self.dmaps[dname]
        else:
            print("No such datamap {} in this Map".format(dname))
            return None

    def _fullname(self, dname):
        return add_suffix(self.mname, dname, separator=default_data_separator)

    def _has_datamap(self, dname):
        return dname in self.dmaps.keys()

    def _regrid_xydatamaps(self):
        if not check._check_ifnD([self.X], ndim=2):
            print("WARNING: regridding X, Y and datamaps into 2D arrays")
            newextent, newX, newY = transform.regrid_XY(self.X, self.Y)
            for dname in self.dmaps.keys():
                self.dmaps[dname].data = transform.regrid_Z(self.X, self.Y,
                                                           self.dmaps[dname].data,
                                                           newX, newY,
                                                           fill_value=self._fill_value,
                                                           method=self._method)
                self.dmaps[dname].edata = transform.regrid_Z(self.X, self.Y,
                                                            self.dmaps[dname].edata,
                                                            newX, newY,
                                                            fill_value=self._fill_value,
                                                            method=self._method)
            # Finally getting the new X and Y
            self.X, self.Y = newX, newY
            self.shape = self.X.shape

    def _reshape_datamaps(self):
        """Reshape all datamaps following X,Y shape
        """
        for dname in self.dmap.keys():
            self.dmaps[dname].reshape(self.shape)

    def attach_datamap(self, datamap):
        """Attach a DataMap to this Map

        Args:
            datamap: a DataMap
        """
        if self._check_datamap(datamap):
            datamap._reshape_datamap(self.shape)
            self.dmaps[datamap.dname] = datamap
            print("INFO: Attaching datamap {0} of type {1} (unit = {2})".format(
                      datamap.dname, datamap.flag, datamap.dunit))
        else:
            print("WARNING[attach_datamap]: could not attach datamap")

    def _check_datamap(self, datamap):
        """Check consistency of data
        """
        # Main loop on the names of the dmaps
        arrays_to_check = [datamap.data.ravel()]
        if datamap.edata is not None:
            arrays_to_check.append(datamap.edata.ravel())

        # First checking that the data are arrays
        if not check._check_ifarrays(arrays_to_check):
            print("ERROR[check_datamap]: input maps not all arrays")
            return False

        # Then checking that they are consistent with X, Y
        arrays_to_check.insert(0, self.X.ravel())
        if not check._check_consistency_sizes(arrays_to_check):
            print("ERROR[check_datamap]: input datamap does not "
                  "have the same size than input grid (X, Y)")
            return False

        return True

    def align_axes(self, galaxy):
        """Align all axes using X and Y as input
        """
        self.align_xy_lineofnodes(galaxy)
        self.align_xy_bar(galaxy)
        self.align_xy_deproj_bar(galaxy)

    @property
    def XY_extent(self):
        return [np.min(self.X), np.max(self.X),
                np.min(self.Y), np.max(self.Y)]

    @property
    def _R(self):
        return np.sqrt(self.X**2 + self.Y**2)

    # Setting up NE direct or not
    @property
    def NE_direct(self) :
        return self.__NE_direct

    @NE_direct.setter
    def NE_direct(self, NE_direct) :
        self.__NE_direct = NE_direct
        self._mat_direct = np.where(NE_direct,
                                    transform.set_stretchmatrix(),
                                    transform.set_reverseXmatrix())
    # Setting up North-East to the top
    @property
    def alpha_North(self) :
        return self.__alpha_North

    @alpha_North.setter
    def alpha_North(self, alpha_North) :
        """Initialise the parameters in the disc structure for alpha_North angles
        in degrees and radian, as well as the associated transformation matrix

        Input
        -----
        alpha_North: angle in degrees for the PA of the North direction
        """
        self.__alpha_North = alpha_North
        self.__alpha_North_rad = deg2rad(alpha_North)
        self._mat_NE = self._mat_direct @ transform.set_rotmatrix(self.__alpha_North_rad)

    def _get_angle_from_PA(self, PA):
        """Provide a way to get the angle within the original
        frame of a certain axis with a given PA
        Args:
            PA: float
                PA of axis with respect to North

        Returns:
            The angle in the original frame
        """
        return PA + self.alpha_North * np.where(self.NE_direct, 1., -1.)

    def align_xy_NorthEast(self) :
        """Get North to the top and East on the left
        """
        self.X_NE, self.Y_NE = self.rotate(matrix=self._mat_NE)

    def align_xy_lineofnodes(self, galaxy) :
        """Set the Line of Nodes (defined by its Position Angle, angle from the North
        going counter-clockwise) as the positive X axis
        """
        self._mat_lon_NE = galaxy._mat_lon.dot(self._mat_NE)
        self.X_lon, self.Y_lon = self.rotate(matrix=self._mat_lon_NE)

    def deproject(self, galaxy):
        """Deproject X,Y around the line of nodes using the inclination
        """
        self.X_londep, self.Y_londep = self.rotate(matrix=galaxy._mat_inc,
                                             X=self.X_lon, Y=self.Y_lon)

    def align_xy_bar(self, galaxy) :
        """Set the bar (defined by its Position Angle, angle from the North
        going counter-clockwise) as the positive X axis
        """
        self.X_bar, self.Y_bar = self.rotate(matrix=galaxy._mat_bar @ self._mat_NE)

    def align_xy_deproj_bar(self, galaxy) :
        """Set the bar (defined by its Position Angle, angle from the North
        going counter-clockwise) as the positive X axis after deprojection
        """
        self._mat_deproj_bar = galaxy._mat_bardep @ galaxy._mat_inc @ galaxy._mat_lon @ self._mat_NE
        self.X_bardep, self.Y_bardep = self.rotate(matrix=self._mat_deproj_bar)

        ## Mirroring the coordinates
        self.X_mirror, self.Y_mirror = self.rotate(matrix=np.linalg.inv(self._mat_deproj_bar),
                                                   X=self.X_bardep, Y=-self.Y_bardep)

    def rotate(self, **kwargs):
        """Uses the rotate function from transform.py with a default
        X,Y set of arrays using self.X and self.Y

        Parameters
        ----------
        **kwargs: set of arguments, see transform.rotate
            Includes X, Y, matrix

        Returns:
        The rotated arrays Xrot, Yrot
        """
        X = kwargs.pop("X", self.X)
        Y = kwargs.pop("Y", self.Y)
        return transform.rotate_vectors(X, Y, **kwargs)

    def intmap_to_densitymap(self, dname, galaxy):
        """Change intensity into density quantity
        by dividing by the XYunit**2

        Args:
            dname (str): name of the datamap
            galaxy (Galaxy):

        Does:
            attach a new map with the normalisation

        """
        if not self._has_datamap(dname):
            return

        dmap = self.dmaps[dname]
        # Test if the map is a density one using the type
        if not _is_flag_density(dmap.flag):
            scalepc2 = galaxy.pc_per_xyunit(self.XYunit) ** 2
            newdata = dmap.data / scalepc2
            if dmap.edata is not None:
                newedata = dmap.edata / scalepc2
            else:
                newedata = None
            newdtype = _add_density_suffix(dmap.dtype)
            newdunit = dmap.dunit / self.XYunit**2
            newflag = _add_density_prefix(dmap.flag)
            newdname = _add_density_prefix(dmap.dname)
            self.add_datamap(dname=newdname, data=newdata, edata=newedata,
                             dunit=newdunit, dtype=newdtype,
                             flag=newflag, order=dmap.order,
                             comment="Renormalised density")
            return newdname
        else:
            return dname

    def deproject_velocities(self, dname, inclin=90.0):
        """Deproject Velocity map if it exists

        Parameters
        ----------
        dname: str
            Name of the datamap to deproject
        inclin: float [90]
            Inclination in degrees
        """

        if dname in self.dmaps:
            self.dmaps[dname].deproject_velocities(inclin=inclin)
        else:
            print("ERROR: no such data name in this Map")

class Profile(object):
    """A Profile is a set of DataProfiles associated via the same R profile.
    It is used to describe radial dprofiles e.g., rotation curves.

    Attributes
    ----------
    R: numpy array [None]
        Input location radii
    data: numpy array [None]
        Input values
    edata: numpy array [None]
        Uncertainties
    order: int [0]
        Order of the velocity moment. -1 for 'others' (e.g., X, Y)
    """
    def __init__(self, R=None, ref_size=None,
                 pname=None, ptype="", **kwargs):
        """
        Args:
            R (numpy array): radii
            pname (str): name of the profile [None]
            ptype (str): type of the profile

            **kwargs:
                data (array): input data
                edata (array): uncertainties
                dname (str): name of the data
                ptype (str): type of the dataprofile
                flag (str): flag for the dataprofile
                order (int): order
                comment (str): comment to be attached [""]
        """
        # Empty dictionary for the moments
        self.dprofiles = AttrDict()

        # See if a dataprofile is provided
        self.Runit = kwargs.pop("Runit", dict_units['R'])
        self.pixel_scale = kwargs.pop("pixel_scale", 1.)

        # Get the list of suffixes which will be used to attach datasets
        dict_dprofs, prof_kwargs = analyse_data_kwargs(**kwargs)

        # First getting the shape of the data
        if ref_size is not None:
            self.size = ref_size
        elif R is not None:
            self.size = R.size
        else:
            # Look for the first data which exists
            data = None
            for dname in dict_dprofs.keys():
                if "data" in dict_dprofs[dname]:
                    data = dict_dprofs[dname]['data']
                    break
            if check._check_ifarrays([data]):
                self.size = data.size
            else:
                print("ERROR: no reference shape is provided "
                      "(via R, data or ref_size) - Aborting")
                return

        # New step in R when provided
        Rfinestep = kwargs.pop("Rfinestep", 0)
        self._init_R(R)

        # Filling value
        self._fill_value = kwargs.pop("fill_value", 'nan')
        # Method
        self._method = kwargs.pop("method", "linear")

        # Comment for Profile
        self.comment = kwargs.pop("comment", "")
        # Name of Profile
        self.pname = pname
        self.ptype = ptype
        self.overwrite = kwargs.pop("overwrite", False)
        self._force_dtype = kwargs.pop("force_dtype", False)

        # Add each datamap one by one
        for dname, dprof_kwargs in dict_dprofs.items():
            if 'dname' not in dprof_kwargs:
                dprof_kwargs['dname'] = dname
            self.add_dataprofile(**dprof_kwargs)

        if Rfinestep > 0:
            self.interpolate(newstep=Rfinestep)

    @property
    def _get_pixel_scale(self):
        return (1. * u.pixel).to(self.Runit, equivalencies=self.eq_pscale).value

    def _convert_to_runit(self):
        """Convert XYunit into the default one
        a priori arcseconds.
        """
        self.R *= self.Runit_per_pixel
        # Update the unit
        self.Runit = dict_units['R']

    @property
    def Runit_per_pixel(self):
        return (1. * self.Runit).to(dict_units['R'],
                                 self.eq_pscale).value
    @property
    def eq_pscale(self):
        return u.pixel_scale(self.pixel_scale * self.Runit / u.pixel)

    def _init_R(self, R):
        """Initialise Rin

        Args:
            R: numpy array
        """
        # Define the grid in case Rin
        # If it is the case, using the reference profile
        if R is None:
            self.R = np.arange(self.size, dtype=default_float)
        else:
            self.R = R.ravel()
        self._convert_to_runit()

    @property
    def ndataprofs(self):
        return len(self.dprofiles)

    def _fullname(self, dname):
        return add_suffix(self.pname, dname, separator=default_data_separator)

    def _has_dataprofile(self, dname):
        return dname in self.dprofiles.keys()

    def _get_dataprofile(self, dname=None, order=None):
        """Get the dataprofile if it exists, and
        check the order

        Args:
            dname:
            order:

        Returns:

        """
        if dname is None:
            if order is None:
                # then just get the first profile
                dname = list(self.dprofiles.keys())[0]
            else:
                # Then get the first profile of right order
                for key in self.dprofiles.keys():
                    if self.dprofiles[key].order == order:
                        dname = key
                        break

        if self._has_dataprofile(dname):
            return self.dprofiles[dname]
        else:
            print("No such dataprofile {} in this Map".format(dname))
            return None

    def attach_dataprofile(self, dataprofile):
        """Attach a DataProfile to this Profile

        Args:
            dataprofile: DataProfile to attach
        """
        if self._check_dprofiles(dataprofile):
            self.dprofiles[dataprofile.dname] = dataprofile

    def add_dataprofile(self, **kwargs):
        """Attach a new Profile to the present Set.

        Args:
            data: 1d array
            order: int
            edata: 1d array
            dname: str
            dtype: str
            flag: str
            dunit: astropy unit

        """
        data = kwargs.pop("data", None)
        if data is None:
            print("ERROR[add_dataprofile]: data is None - Aborting")
            return

        # Input dname to define the data. If none, define using the counter
        dname = kwargs.pop("dname", None)
        if dname is None or dname=="":
            dname = "{0}{1:02d}".format(self.pname, self.ndataprofs+1)
        if dname[0] == default_data_separator:
            dname = dname[1:]

        if self._has_dataprofile(dname) and not overwrite:
            print("WARNING[add_dataprofile]: data profile {} already exists "
                  "- Aborting".format(dname))
            print("WARNING[add_dataprofile]: use overwrite option to force.")
            return

        # Check if we wish to force the dtype / dunit
        force_dtype = kwargs.pop("force_dtype", self._force_dtype)
        if force_dtype:
            dtype = kwargs.get("dtype", None)
            dmap_info = get_dmap_info_from_dtype(dtype, dname)
            # Transfer
            for key, value in dmap_info.items():
                kwargs[key] = value

        self.attach_dataprofile(DataProfile(data=data, dname=dname, **kwargs))

    def __getattr__(self, dname):
        for suffix in default_data_names:
            if dname.startswith(suffix):
                for profname in self.dprofiles.keys():
                    if profname in dname:
                        basename = remove_suffix(dname, profname)
                        return getattr(self.dprofiles[profname], basename)
        raise AttributeError("'Profile' object has no attribute {}".format(dname))

    def __dir__(self, list_names=default_data_names):
        return  super().__dir__() + [add_suffix(attr, prof) for item in list_names
                for prof in self.dprofiles.keys() for attr in self.dprofiles[prof].__dir__()
                if attr.startswith(item)]

    def _check_dprofiles(self, dataprofile):
        """Check consistency of dataprofile
        by comparing with self.Rin

        Args
            dataprofile: DataProfile
        """
        # Putting everything in 1D
        ref_array = self.R

        data = dataprofile.data
        edata = dataprofile.edata
        arrays_to_check = [data.ravel()]
        if edata is not None:
            arrays_to_check.append(edata.ravel())

        # Checking if the data are 1D arrays
        if not check._check_ifarrays(arrays_to_check):
            print("ERROR: input profile not all arrays")
            return False

        # Check that they all have the same size
        arrays_to_check.insert(0, ref_array)
        if not check._check_consistency_sizes(arrays_to_check):
            print("ERROR: input profile does not the same size "
                  "than input radial grid (R)")
            return False

        return True

    def interpolate(self, dname, step=1.0, suffix="fine", overwrite=False):
        """Provide interpolated profile

        Args:
            stepR: float [1.0]
            suffix: str [""]
            overwrite: bool [False]

        Returns:

        """
        if step <= 0:
            print("ERROR[interpolate]: new step is <= 0 - Aborting")
            return

        # Getting the data
        if not self._has_dataprofile(dname):
            print("ERROR[interpolate]: no such dataprofile "
                  "with name {}".format(dname))
            return

        if hasattr(self.dprofiles[dname], add_suffix("R", suffix)):
            if overwrite:
                print("WARNING: overwriting existing interpolated profile")
            else:
                print("ERROR[interpolate]: interpolated profile exists. "
                      "Use 'overwrite' to update.")
                return

        Rfine, dfine, edfine = transform.interpolate_profile(self.R,
                                                             self.dprofiles[dname].data,
                                                             self.dprofiles[dname].edata,
                                                             step=step)
        setattr(self.dprofiles[dname], add_suffix("R", suffix), Rfine)
        setattr(self.dprofiles[dname], add_suffix("data", suffix), dfine)
        setattr(self.dprofiles[dname], add_suffix("edata", suffix), edfine)

class Slicing(object):
    """Provides a way to slice a 2D field. This class just
    computes the slits positions for further usage.
    """
    def __init__(self, yextent=[-10.,10.], yin=None, slit_width=1.0, nslits=None):
        """Initialise the Slice by computing the number of slits and
        their positions (defined by the axis 'y').

        Args:
            yextent: list of 2 floats
                [ymin, ymax]
            yin: numpy array
                input y position
            slit_width: float
                Width of the slit
            nslits: int
                Number of slits. This is optional if a range or input yin
                is provided.
        """

        # First deriving the range. Priority is on yin
        if yin is not None:
            yextent = [np.min(yin), np.max(yin)]
        # First deriving the number of slits prioritising nslits

        Dy = np.abs(yextent[1] - yextent[0])
        if nslits is None:
            self.nslits = np.int(Dy / slit_width + 1.0)
            ye2 = (Dy - self.nslits * slit_width) / 2.
            # Adding left-over on both sides equally
            yextent = [yextent[0] - ye2, yextent[1] + ye2]
        else:
            self.nslits = nslits
            slit_width = Dy / self.nslits

        self.width = slit_width
        sw2 = slit_width / 2.
        self.ycentres = np.linspace(yextent[0] + sw2, yextent[1] - sw2, self.nslits)
        self.yedges = np.linspace(yextent[0], yextent[1], self.nslits+1)
        self.yiter = np.arange(self.nslits)

        @property
        def yextent(self):
            return [self.yedges[0], self.yedges[-1]]

        @property
        def slice_width(self):
            return np.abs(self.yedges[-1] - self.yedges[0])

def match_datamaps(map1, map2=None, dname1=None, dname2=None,
                   odname1=None, odname2=None, PAnodes=0.):
    """Aligning two datamaps

    Args:
        map1 (Map): input Map
        map2 (Map): second input Map. If None, use the first one.
        dmap1_name (str): name of input datamap 1
        dmap2_name (str): name of input datamap 2
        omap1_name (str): name of output datamap 1
        omap2_name (str): name of output datamap 2

    Returns:
        New Map with matched datamaps
    """
    if map2 is None:
        map2 = map1

    # Get the datamaps
    dmap1 = map1._get_datamap(dname1)
    dmap2 = map2._get_datamap(dname2)
    if dmap1 is None or dmap2 is None:
        return None

    # Determine the new grid
    XYextent = get_extent(map1.X_lon, map1.Y_lon)
    newstep = guess_stepxy(map1.X_lon, map1.Y_lon)
    Xn, Yn = np.meshgrid(cover_linspace(XYextent[0], XYextent[1], newstep),
                         cover_linspace(XYextent[2], XYextent[3], newstep))

    # Regrid
    new_data1 = transform.regrid_Z(map1.X_lon, map1.Y_lon, dmap1.data, Xn, Yn)
    new_edata1 = transform.regrid_Z(map1.X_lon, map1.Y_lon, dmap1.edata, Xn, Yn)
    new_data2 = transform.regrid_Z(map2.X_lon, map2.Y_lon, dmap2.data, Xn, Yn)
    new_edata2 = transform.regrid_Z(map2.X_lon, map2.Y_lon, dmap2.edata, Xn, Yn)

    # And re-attach to a regrided mass map
    omname1 = add_suffix(map1.mname, remap_suffix, separator="")
    if odname1 is None:
        odname1 = add_suffix(dmap1.dname, remap_suffix, separator="")
    if odname2 is None:
        odname2 = add_suffix(dmap2.dname, remap_suffix, separator="")
    mtype1 = add_prefix(map1.mtype, remap_suffix, separator="")
    dtype1 = dmap1.dtype.lower()
    dtype2 = dmap2.dtype.lower()

    # Creating the new Map
    print("INFO[match_datamaps]: Creating the first map {0} and "
          "attaching first datamap {1}".format(omname1, dname1))
    newMap = Map(mname=omname1, data=new_data1, edata=new_edata1, order=0,
                 mtype=mtype1, X=Xn, Y=Yn, dtype=dtype1, flag=dmap1.flag,
                 dunit=dmap1.dunit, dname=odname1, alpha_north=-90.0-PAnodes)

    # Adding the second datamap
    print("INFO[match_datamaps]: attaching the datamap {0} to map {1}".format(
           dname2, omname1))
    newMap.add_datamap(data=new_data2, order=dmap2.order, edata=new_edata2,
                    dname=odname2, flag=dmap2.flag, dtype=dtype2,
                    dunit=dmap2.dunit)

    return newMap


import os
import sys
import csv
import numpy as np
import logging
import yaml
import pytz
from pathlib import Path
from datetime import datetime, timedelta
from matplotlib import pyplot as plt
import mne

# How wide of a buffer around the crop do we want?
# 1 second is enough with .5s epochs
BUFFER_SECONDS = 1

HIGHPASS_ARTIFACT = 0.5
HIGHPASS_MMN = 1
LOWPASS_MMN = 35
HIGHPASS_ABR = 100
LOWPASS_ABR = 3000

# Event IDs
UNKNOWN = 1
STANDARD = 2
DEVIANT = 3
EVENT_COLORS = {
    UNKNOWN: "blue",
    STANDARD: "green",
    DEVIANT: "red"
}

class BDFWithMetadata():
    def __init__(self, path, kind, force=False, is_2013I=False, no_reference=False, reference_o1=False, reference_o2=False, no_notch=False, no_crop=False):
        self.script_dir = sys.path[0]
        self.kind = kind
        self.is_2013I = is_2013I
        self.no_reference = no_reference
        # Both is the default, so set them to false
        if reference_o1 and reference_o2:
            reference_o1 = False
            reference_o2 = False
        self.reference_o1 = reference_o1
        self.reference_o2 = reference_o2
        self.no_notch = no_notch
        self.no_crop = no_crop

        # Determine if source path is in the standard /study/thukdam/raw-data/subjects location or not
        p = Path(path).resolve()
        self.source_path = str(p)
        parts = p.parts
        if "raw-data" in parts and "subjects" in parts:
            # If "raw-data", "subjects" is in path, replace "raw-data" with "analyses", "eeg"
            dest = list(parts)
            raw_index = dest.index('raw-data')
            dest[raw_index] = "analyses"
            dest.insert(raw_index+1, "eeg_artifacts")

            # Kind is MMN or ABR, but now we want to allow for O1/O2 only referencing
            subfolder = kind
            if self.reference_o1:
                subfolder += "-o1"
            if self.reference_o2:
                subfolder += "-o2"
            dest[raw_index+2] = subfolder
            dest.remove("biosemi")
            dest[-1] = dest[-1].replace(".bdf", "")

            artifact_path = Path(os.path.join(*dest))
            dest[raw_index+1] = "eeg_plots"
            plot_path = Path(os.path.join(*dest))

            dest[raw_index+1] = "eeg_statistics"
            statistics_path = Path(os.path.join(*dest))
        else:
            artifact_path = p
            plot_path = p
            statistics_path = p
            if not force:
                # DIE unless they forced to save in current dir with a flag
                logging.critical("Data file not stored in expected raw-data/subjects directory, please run with --force to save masks and plots and stuff to current directory")
                sys.exit(1)

        output_dir = p.parent
        self.artifact_path = str(artifact_path)
        self.plot_path = str(plot_path)
        self.statistics_path = str(statistics_path)

        logging.info(f"Saving artifacts to {self.artifact_path}")
        logging.info(f"Saving plots to {self.plot_path}")
        logging.info(f"Saving statistics to {self.statistics_path}")

        # Create directories if they don't exist
        artifact_path.parent.mkdir(parents=True, exist_ok=True)
        plot_path.parent.mkdir(parents=True, exist_ok=True)
        statistics_path.parent.mkdir(parents=True, exist_ok=True)

        self.highpass_artifact = HIGHPASS_ARTIFACT

        if self.is_mmn():
            self.highpass = HIGHPASS_MMN
            self.lowpass = LOWPASS_MMN
            self.event_id = {
                "Unknown": UNKNOWN,
                "Standard": STANDARD,
                "Deviant": DEVIANT
            }
        else:
            self.highpass = HIGHPASS_ABR
            self.lowpass = LOWPASS_ABR
            self.event_id = {
                "Unknown": UNKNOWN,
            }

    def is_standard_frequencies(self):
        if self.is_mmn():
            return self.highpass == HIGHPASS_MMN and self.lowpass == LOWPASS_MMN
        else:
            return self.highpass == HIGHPASS_ABR and self.lowpass == LOWPASS_ABR

    def is_mmn(self):
        return self.kind == "mmn"

    def load(self):
        # Do we have existing metadata?
        self.load_existing_metadata()
        self.load_existing_events()

        self.load_file(self.source_path)

    def strip_reference_electrode(self, path):
        return path.replace("-o1", "").replace("-o2", "")

    def load_existing_metadata(self):
        metadata = self.artifact_metadata_file()

        # If we can't find the file at the default path, try stripping out -o1 or -o2
        # to get the "previous" mask with both reference electrodes
        if not os.path.exists(metadata):
            metadata = self.strip_reference_electrode(metadata)

        if os.path.exists(metadata):
            with open(metadata) as file:
                data = yaml.load(file, Loader=yaml.FullLoader)
                self.tstart_seconds = data['tstart_seconds']
                self.tstop_seconds = data['tstop_seconds']
                if 'highpass' in data:
                    self.highpass = data['highpass']
                if 'lowpass' in data:
                    self.lowpass = data['lowpass']
                logging.info(f"Loaded existing start {self.tstart_seconds} and end {self.tstop_seconds} from {metadata}, frequencies are {self.highpass}Hz to {self.lowpass}Hz")
        else:
            self.tstart_seconds = None
            self.tstop_seconds = None

    def load_existing_events(self):
        e = self.events_file()

        # If we can't find the file at the default path, try stripping out -o1 or -o2
        # to get the "previous" events file from the artifact rejection with both reference electrodes
        if not os.path.exists(e):
            e = self.strip_reference_electrode(e)

        if os.path.exists(e):
            self.events = np.load(e)
            logging.info(f"Loaded {len(self.events)} from events numpy file {e}")
        else:
            self.events = []

    def locate_events(self, expected_events, expected_duration, kind):
        """
        Locate event chunks in file that match the given duration.

        expected_events: Integer number of events we want to find
        expected_duration: Duration in seconds that we want to find them in
        kind: User-visible sort of events we're looking for
        """
        sfreq = self.raw.info['sfreq']
        raw_events = mne.find_events(self.raw)

        skipped_events = 0
        looking = True

        logging.info(f"Found {len(raw_events)} total events in file.")
        if len(raw_events) < expected_events:
            logging.fatal(f"Not enough events to find {expected_events}, exiting!")
            sys.exit(1)
        while looking and len(raw_events) - skipped_events >= expected_events:
            tstart = (raw_events[skipped_events,0] - (sfreq * BUFFER_SECONDS))
            self.tstart_seconds = tstart / sfreq
            tstop = (raw_events[skipped_events+expected_events-1,0] + (sfreq * BUFFER_SECONDS))
            self.tstop_seconds = tstop / sfreq
            duration_seconds = self.tstop_seconds - self.tstart_seconds

            if duration_seconds > expected_duration - BUFFER_SECONDS*2 and \
                duration_seconds < expected_duration + BUFFER_SECONDS*2:
                logging.info(f"Found events at {tstart} with duration {duration_seconds} after skipping {skipped_events}")
                looking = False
            else:
                logging.debug(f"Skipped {skipped_events}, at {tstart} with {duration_seconds} (looking for {expected_events})")
                skipped_events += 1

        if looking:
            # Sorry if this sucks in actual use, bit of a rush to get this all working
            logging.warning(f"Could not find {expected_events} {kind} events automatically, skipped {skipped_events} while trying.")
            logging.warning(f"Please scroll and find start and stop time in seconds manually!")
            # Temporarily set our events to the full list for plotting
            self.events = raw_events.copy()
            self.plot(False)
            self.tstart_seconds = float(input(f"Enter {kind} start time (in seconds): "))
            self.tstop_seconds = float(input(f"Enter {kind} stop time (in seconds): "))

        # Crop to those seconds
        self.raw.crop(tmin=self.tstart_seconds, tmax=self.tstop_seconds)
        self.tstart = self.tstart_seconds * sfreq
        self.tstop = self.tstop_seconds * sfreq

        # Truncate the events list to the ones we wanted,
        # keeping in mind we have to start at the index after the start
        index = np.searchsorted(raw_events[:,0], self.tstart)
        self.events = raw_events[index:index+expected_events].copy()

    def load_file(self, raw_file):
        self.raw = mne.io.read_raw_bdf(raw_file)

        # TODO: Original script does weird event deletion, with this comment:
        """
        On the South computer, it appears that there are many short (2 or 3
        sample duration) events that are probably due to a problem in the MMN
        .WAV files or noise / dropouts / flakiness in the cables connecting the
        field laptop audio-out jack to the input on the BioSemi system. Whatever
        the cause, they're bogus and they screw up analysis. So we'll just delete
        them.
        """

        first_run = True

        if self.tstart_seconds and len(self.events) > 0:
            # If we already have information, use that
            self.raw.crop(tmin=self.tstart_seconds, tmax=self.tstop_seconds)
            first_run = False
        elif not self.no_crop:
            if self.is_mmn():
                # Crop to the MMN section of the file
                self.locate_events(2000, 1000, self.kind)
            else:
                # Crop to the ABR section of the file
                self.locate_events(4000, 200, self.kind)

        self.raw.load_data()
        # Rename channels in raw based on actual electrode names
        # This is based on FMed_Chanlocs_6channels.ced
        # EXG2 used to be called mr and EXG3 was ml,
        # they are renamed to O1 and O2
        # so that when we load the electrode montage below it matches

        if self.raw.info['nchan'] == 17:
            # Unclear why, but in the original files, the 6 channels are duplicated
            # MNE appends a number to differentiate the dupes
            self.raw.rename_channels({'EXG1-0': 'Cz', 'EXG2-0': 'O1', 'EXG3-0': 'O2', 'EXG4-0': 'Fz', 'EXG5-0': 'Pz', 'EXG6-0': 'T8'})
        else:
            self.raw.rename_channels({'EXG1': 'Cz', 'EXG2': 'O1', 'EXG3': 'O2', 'EXG4': 'Fz', 'EXG5': 'Pz', 'EXG6': 'T8'})

        # Reference electrodes on mastoids
        if self.no_reference:
            logging.warning("Not referencing mastoids, raw view")
        else:
            if self.reference_o1:
                logging.warning("Referencing only O1")
                self.raw.set_eeg_reference(['O1'])
            elif self.reference_o2:
                logging.warning("Referencing only O2")
                self.raw.set_eeg_reference(['O2'])
            else:
                self.raw.set_eeg_reference(['O1', 'O2'])

        # Try to hack in some electrode location information into the raw.info
        montage = mne.channels.make_standard_montage('biosemi16')
        self.raw.set_montage(montage, raise_if_subset=False)

        # Notch out the India power frequency unless told not to
        if self.no_notch:
            logging.info("Not notch filtering at 50Hz")
        else:
            logging.info("Notch filtering at 50Hz")
            self.raw.notch_filter(np.arange(50, 251, 50))

        if self.no_crop:
            logging.warning("Not cropping, so not doing any artifact or event loading")
            return
        elif first_run:
            if self.is_mmn():
                # Figure out if tones are same or deviant
                self.load_event_tones_for_mmn()
            # We don't need to do any hacking of events for ABR

            # Now we automatically save out the cropping and events metadata
            self.save_metadata()
            self.save_events()

        # If previous annotations exist, read them
        self.load_annotations()

    def save_metadata(self):
        data = {
            'tstart_seconds': int(self.tstart_seconds),
            'tstop_seconds': int(self.tstop_seconds),
            'source_path': self.source_path,
            'highpass_artifact': self.highpass_artifact,
            'highpass': self.highpass,
            'lowpass': self.lowpass,
        }
        with open(self.artifact_metadata_file(), 'w') as file:
            yaml.dump(data, file)

    def save_events(self):
        np.save(self.events_file(), self.events)

    def load_annotations(self):
        mask_path = self.artifact_mask_file()

        # If we can't find the file at the default path, try stripping out -o1 or -o2
        # to get the "previous" mask file from the previous artifact rejection process
        if not os.path.exists(mask_path):
            mask_path = self.strip_reference_electrode(mask_path)

        if os.path.exists(mask_path):
            logging.info(f"Loading existing artifact annotations from {mask_path}")
            a = mne.read_annotations(mask_path)
            self.raw.set_annotations(a)

    def load_event_tones_for_mmn(self):
        logging.info(f"Determining MMN event types")
        # Now we need to load the right event tones and paste them into the event array
        mmnToneDir = os.path.join(self.script_dir, 'MMN_tone_sequences')
        is2013Initial = self.is_2013I

        if is2013Initial:
            # 2 seconds to account for user accepting MMN .WAV file to be played
            # 303 seconds to play "silence" .WAV file.
            # 1 second to avoid interference between file playback routines
            # 12 seconds from start of MMN .WAV file to first tone being played
            secondsFromScriptStartToFirstTone = 2 + 303 + 1 + 12
            mmnToneFileStart = os.path.join(mmnToneDir, 'south', 'MMN_roving_with_trigger_dpdb02_seed_10')
            mmnToneFileEnd = '_31-May-2014_tone_sequence.txt'
        else:
            # 2 seconds to account for user accepting MMN .WAV file to be played
            # 300 seconds for 300 second "silence" pause
            # 10 seconds from start of MMN .WAV file to first tone being played
            secondsFromScriptStartToFirstTone = 2 + 300 + 10;
            mmnToneFileStart = os.path.join(mmnToneDir, 'north', 'MMN_roving_with_trigger_dpdb01_seed_10')
            mmnToneFileEnd = '_21-Dec-2012_tone_sequence.txt'

        # NOTE: This code matches what the original Matlab script does,
        # but note that we're assuming UTC which feels... strange.
        meas_date = self.raw.info['meas_date'][0]
        logging.info(f'Measured date string in BDF file is {meas_date}')
        recordedDate = datetime.fromtimestamp(meas_date, pytz.timezone("UTC"))

        logging.info(f'Recorded date is {recordedDate}')
        actualStart = recordedDate + timedelta(seconds=self.tstart_seconds - secondsFromScriptStartToFirstTone)
        doy = actualStart.timetuple().tm_yday

        logging.info(f'Script start day of year = {doy}')
        noon = actualStart.replace(hour=12, minute=0, second=0)
        if abs(noon - actualStart).seconds < 600:
            logging.warning(f"WARNING: start time {actualStart} is close to noon, so the event tone discovery may be wrong")

        # We can now guess which tone sequence .TXT file to use for assigning
        # tone IDs to events in the .BDF file.
        logging.info(f'Script actual start is {actualStart}')

        # This determination is based on day of the year and time of day:
        if (actualStart.hour >= 12):
            daySegment = 1
        else:
            daySegment = 0
        dayEven = doy % 2
        whichSeq = 2 * dayEven + daySegment
        mmnToneFileName = f"{mmnToneFileStart}{whichSeq}{mmnToneFileEnd}"
        logging.info(f"Loading tones from {mmnToneFileName}")

        with open(mmnToneFileName) as csvfile:
            reader = csv.reader(csvfile)
            tones = next(reader)
            # NOTE: there are way more than 2000 entries because of... legacy reasons. IGNORE

        # Finally, we know enough to repair the events in the raw data
        # and mark them same or deviant
        numSameEvents = 0
        numDeviantEvents = 0
        for i in range(1, len(self.events)):
            if (tones[i] == tones[i - 1]):
                self.events[i,2] = STANDARD
                numSameEvents += 1
            else:
                self.events[i,2] = DEVIANT
                numDeviantEvents += 1
        logging.info(f"Determined {numSameEvents} same events and {numDeviantEvents} deviant events")

    def artifact_mask_file(self):
        return self.artifact_path + f".{self.kind}_artifact_mask.csv"

    def artifact_metadata_file(self):
        return self.artifact_path + f".{self.kind}_artifact_metadata.yaml"

    def events_file(self):
        return self.artifact_path + f".{self.kind}_events.npy"

    def plot_output_path(self, name):
        return self.plot_path + f".{self.kind}_{name}.png"

    def plot(self, block=True, display_huge=False, no_events=False):
        if display_huge:
            if self.is_mmn:
                duration = 1000.0
            else:
                duration = 200.0
            scalings = dict(eeg=150e-6)
            events = None
        else:
            duration = 5.0
            scalings = dict(eeg=50e-6)
            events = self.events

        if no_events:
            events = None

        order = [1, 2, 3, 0, 4, 5]
        n_channels = 7

        self.raw.plot(
            block=block,
            n_channels=n_channels,
            remove_dc=True,
            events=events,
            event_color=EVENT_COLORS,
            highpass=self.highpass_artifact,
            duration=duration,
            order=order,
            scalings=scalings)

    def artifact_rejection(self, display_huge=False, no_events=False):
        logging.info("View loaded. Ready for artifact rejection! Press 'a' to start, add a label, and then drag on the graph. Close the view window to continue.")
        self.plot(display_huge=display_huge, no_events=no_events)

        mask_path = self.artifact_mask_file()

        # Save the annotations
        if len(self.raw.annotations) > 0:
            self.raw.annotations.save(mask_path)

    def build_epochs(self):
        # Actually do the real final filtering (happens in-place)
        self.raw.filter(l_freq=self.highpass, h_freq=self.lowpass, fir_design='firwin')

        # Epoching...
        picks = ['Cz', 'Fz', 'Pz', 'T8']
        if self.is_mmn():
            tmin, tmax = -0.1, 0.4
        else:
            tmin, tmax = -0.002, 0.010

        epochs_params = dict(events=self.events, event_id=self.event_id,
                            tmin=tmin, tmax=tmax,
                            picks=picks, reject=None, flat=None)

        self.epochs = mne.Epochs(self.raw, **epochs_params)
        return self.epochs

    def save_figure(self, fig, name, force_name=False):
        if self.is_standard_frequencies() or force_name:
            filename = self.plot_output_path(name)
        else:
            filename = self.plot_output_path(f"{name}_{self.highpass}Hz_to_{self.lowpass}Hz")
        fig.savefig(filename, dpi=300)
        logging.info(f"Saved {name} plot to {filename}")

    def epoch_images(self):
        logging.warning("Plotting epoch image, VERY SLOW")
        fig = self.epochs.plot_image(cmap="YlGnBu_r", group_by=None,
                picks=['Cz', 'Fz', 'T8', 'Pz'], show=False)
        self.save_figure(fig, "epochs")

    def average_output_path(self, name):
        return self.statistics_path + f".{self.kind}-{name}-ave.fif"

    def save_average(self):
        if self.is_mmn():
            deviant = self.epochs["Deviant"].average()
            dfile = self.average_output_path("deviant")
            logging.info(f"Saved evoked averages of deviant events to {dfile}")
            mne.write_evokeds(dfile, deviant)

            standard = self.epochs["Standard"].average()
            sfile = self.average_output_path("standard")
            mne.write_evokeds(sfile, standard)
            logging.info(f"Saved evoked averages of standard events to {sfile}")

        average = self.epochs.average()
        afile = self.average_output_path("all")
        mne.write_evokeds(afile, average)
        logging.info(f"Saved evoked averages to {afile}")

    def epoch_view(self):
        logging.info("Loading epoch viewer...")
        epochs.plot(block=True)

    def psd(self, high_freq):
        # Spectral density is go!
        # https://mne.tools/stable/generated/mne.io.Raw.html#mne.io.Raw.plot_psd
        #fig = self.raw.plot_psd(0, high_freq, average=False, show=False, estimate='power')
        fig = self.raw.plot_psd(0, high_freq, area_mode='std', show=False, n_fft=60000)
        title = f"Power spectral density for {self.kind}"
        self.save_figure(fig, f"psd_to_{high_freq}", True)

    def topo(self):
        epochs = self.epochs

        joint_kwargs = dict(ts_args=dict(time_unit='s'),
                        topomap_args=dict(time_unit='s'),
                        show=False)
        if self.is_mmn():
            deviant = epochs["Deviant"].average()
            standard = epochs["Standard"].average()
            fig1 = deviant.plot_joint(**joint_kwargs)
            self.save_figure(fig1, "deviant_average")
            fig2 = standard.plot_joint(**joint_kwargs)
            self.save_figure(fig2, "standard_average")
        else:
            average = epochs.average()
            fig = average.plot_joint(**joint_kwargs)
            self.save_figure(fig, "epoch_average")


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
import math, copy, csv, sys
from numpy import mat, multiply

class landmarks(object):
    """
    The landmarks class can be used to read, write and manipulate landmarks.  For example...
        lmk = landmarks()                                       # Construct
        lmk = landmarks([ ['label1',1.5,2], ['label2',2.4,5]] ) # Construct from landmarks list
        lmk = landmarks("/path/to/landmarks.lmk")               # Construct from landmarks file
        lmk.Read(path)           # Read landmarks from given path
        labels = lmk.GetLabels() # Get the labels of the landmarks
        lmk.SetLabels(labels)    # Set the labels of the landmarks
        lmk.Write(path)           # Write landmarks to given path
    """

    def __init__(self, inputLandmarks=[['0',0,0,0]], spacing=[1.0,1.0,1.0]):
        self.spacing = spacing
        if type(inputLandmarks) is str:
            self.Read(inputLandmarks, spacing)
        elif type(inputLandmarks) is list:
            self.SetLandmarks(inputLandmarks)
        else:
            raise Exception("Landmarks must be constructed with either a landmaks list or landmarks file path")

    def SetLandmarks(self,landmarkList):
        """
        Sets the landmarks using given landmark list.
        Format of landmarkList must be...
            [[label1,x1,y1,...],...,[labelN,xN,yN,...]]
        ... where x1,y1,...,xN,yN... must be numeric.
        """
        # Make sure that landmarkList is a list
        if not(type(landmarkList) is list): raise Exception("landmarkList must be of type list")

        if len(landmarkList) != 0:
            # Make sure that each landmark in landmarkList is a list
            for landmark in landmarkList:
                if not(type(landmark) is list): raise Exception("Each landmark in landmarkList must be a list.")

            # Make sure dimension of 1st landmark in landmarkList > 0
            dimension = len(landmarkList[0]) - 1
            if dimension < 1: raise Exception("Each landmark in landmarkList must have 1 or more dinemsions")

            for landmark in landmarkList:
                # Make sure each landmark in landmarkList has same dimension
                if len(landmark)-1 != dimension: raise Exception("Each landmark in landmarkList must have same dimension.")

                # Make sure each landmark in landmarkList has numeric cordinates
                for value in landmark[1:]:
                    try:
                        value + 0
                    except TypeError:
                        raise Exception("Coordinates of landmarks must be numeric.")

        self.landmarkList = landmarkList

    def GetLandmarks(self,labelList=[]):
        """
        Returns list of lanmarks in the format
            [[label1,x1,y1,...],...,[labelN,xN,yN,...]]
        """
        if labelList==[]:
            return self.landmarkList[:]
        else:
            outLandmarkList = []
            for label in labelList:
                for landmark in self.landmarkList:
                    if landmark[0] == label: outLandmarkList.append(landmark)
            return outLandmarkList

    def GetPoints(self, labelList=[]):
        """
        Returns list of lanmarks in the format
            [[x1,y1,...],...,[xN,yN,...]]
        """
        if labelList==[]: labelList = self.GetLabels()
        outPointList = []
        for label in labelList:
            for landmark in self.landmarkList:
                if landmark[0] == label: outPointList.append(landmark[1:])
        return outPointList

    def SetLabels(self,labelList):
        """
        Sets the labels of the landmarks where labelList must have format...
             [label1,...,labelN]
        It is necessary that len(labelList) == landmarks.GetNumberOfLandmarks()
        """
        if not(type(labelList) is list): raise Exception("labelList must be of type list")
        if len(labelList) != self.GetNumberOfLandmarks(): raise Exception("Label list must contain same number of labels as this landmarks object.")
        for i in range(0,self.GetNumberOfLandmarks()): self.landmarkList[i][0] = labelList[i]

    def GetLabels(self):
        """
        Returns list of landmark labels in the format...
            [label1,...,labelN]
        """
        labelList = []
        for landmark in self.landmarkList: labelList.append(landmark[0])
        return labelList

    def GetDimensionOfLandmarks(self):
        """
        Returns the dimension of landmarks.
        """
        return len(self.landmarkList[0]) - 1

    def GetNumberOfLandmarks(self):
        """
        Returns the number of landmarks.
        """
        return len(self.landmarkList)

    def GetNumberOfPoints(self):
        """
        Returns the number of points
        """
        return self.GetNumberOfLandmarks()

    def GetSize(self):
        """
        Returns size of bounding box at origin which contains all landmarks
        """
        dimension = self.GetDimensionOfLandmarks()
        size = [0]*dimension
        for lmk in self.landmarkList:
            for i in range(dimension):
                if lmk[i+1] > size[i]: size[i] = lmk[i+1]
        return(size)

    def GetMin(self):
        """
        Returns minimum index of bounding box containing all landmarks
        """
        dimension = self.GetDimensionOfLandmarks()
        minValue = [float('inf')]*dimension
        for lmk in self.landmarkList:
            for i in range(dimension):
                if lmk[i+1] < minValue[i]: minValue[i] = lmk[i+1]
        return minValue

    def GetMax(self):
        """
        Returns maximum index of bounding box containing all landmarks
        """
        dimension = self.GetDimensionOfLandmarks()
        maxValue = [float('-inf')]*dimension
        for lmk in self.landmarkList:
            for i in range(dimension):
                if lmk[i+1] > maxValue[i]: maxValue[i] = lmk[i+1]
        return maxValue

    def Read(self,path,spacing=[1.0,1.0,1.0]):
        """
        Reads landmarks from given path.
        It accepts landmarks files in the following formats
        1. CIS format which usually has extension .lmk
            Landmarks...
            N
            label1
            x1 y1 z1 1 1
            ...
            labelN
            xN yN zN 1 1
        2. DiffeoMap format which usually has extension .txt
            Landmarks...
            N
            label1 x1 y1 z1 ...
            ...
            labelN xN yN zN ...
        """
        self.spacing = spacing

        # Read lines from landmark file into list while triming whitespace and skiping blank line
        landmarkFile = open(path,"r")
        lineList = []
        for line in landmarkFile.readlines():
            if line.strip() != "": lineList.append(line.strip())
        landmarkFile.close()

        # Make sure this is a valid landmarks file by checking that the 1st line begins with "Landmarks"
        if lineList.pop(0)[0:9] != "Landmarks":
            raise Exception(path + " is not a valid landmarks file.")

        # Get number of landmarks from 2nd line of file
        numberOfLandmarks = int(lineList.pop(0))

        # Split all lines in line list into words
        #def StringSplit(string): return string.split()
        lineList = map(str.split,lineList)

        # Read landmarks from line list
        landmarkList = []
        i = 0
        while i < len(lineList):
            if len(lineList[i]) == 1: # For CIS format landmarks
                label = lineList[i][0]
                index = map(float,lineList[i+1][0:3])
                i += 2
            else:                     # For DiffeoMap format landmarks
                label = lineList[i][0]
                index = map(float,lineList[i][1:4])
                i += 1

            point = [a*b for (a,b) in zip(index, spacing)]
            landmark = [label] + point
            landmarkList.append(landmark)
            if len(landmarkList) >= numberOfLandmarks: break

        self.SetLandmarks(landmarkList)

    def Write(self,path):
        """
        Writes landmarks to given path using CIS format.
        """
        # Write header
        landmarkFile = open(path,"w")
        print("Landmarks-1.0",file=landmarkFile)              # 1st line should always be "Landmarks-1.0"
        print(self.GetNumberOfLandmarks(), file=landmarkFile) # 2nd line is always the number of landmarks

        # Write landmarks
        for landmark in self.landmarkList:
            label = landmark[0]
            point = landmark[1:]
            index = [str(a/b) for (a,b) in zip(point, self.spacing)]
            print(label,file=landmarkFile)
            print(" ".join(index + ["1","1"]), file=landmarkFile)

        # Write tail
        # tail = '0\n0\n0\n0,1,0\n0\n"NeuroData"\n0.1,0.9\n"Voxel"\n0,0,0\n' + ",".join(map(str,self.spacing))
        # print(tail, file=landmarkFile)
        landmarkFile.close()

    def WriteCsv(self, path):
        csvFile = open(path, 'w')
        csvWriter = csv.writer(csvFile)
        header = ['label','x','y','z']
        csvWriter.writerows([header] + self.landmarkList)
        csvFile.close()

    def GetDistances(self,otherLandmarks,spacing=[1,1,1]):
        """
        Returns a list of distances between corresponding landmarks of this landmarks object and another landmarks objects.
        """
        if self.GetNumberOfLandmarks() != otherLandmarks.GetNumberOfLandmarks(): raise Exception("Other landmarks must have same number of landmarks as this landmarks object.")

        otherLandmarkList = otherLandmarks.GetLandmarks()
        distanceList = []
        for i in range(0,self.GetNumberOfLandmarks()):
            sumOfSquaredDifferences = 0
            for j in range(0,self.GetDimensionOfLandmarks()):
                sumOfSquaredDifferences += ((self.landmarkList[i][j+1] - otherLandmarkList[i][j+1])*spacing[j])**2

            distanceList.append(math.sqrt(sumOfSquaredDifferences))

        return distanceList

    def Affine(self, affine):
        print(type(self))
        ### TODO generalize to other dimensions
        if (not(type(affine)) is list) or (len(affine) != 12): raise Exception("affine must be a list of length 12.")
        lmkList = []
        labelList = self.GetLabels()
        for i in range(self.GetNumberOfLandmarks()):
            x0 = mat(self.landmarkList[i][1:]).reshape(3,1)
            A = mat(affine[:9]).reshape(3,3)
            b = mat(affine[9:]).reshape(3,1)
            #x1 = A.I*(x0 - b)
            x1 = A*x0 + b
            lmk = [labelList[i]] +  x1.flatten().tolist()[0]
            lmkList.append(lmk)
        return landmarks(lmkList, spacing=self.spacing)

    def Flip(self, size):
        """
        Returns flipped landmarks.
        Useful for converting BrainWorks Landmarks to MRIStudio Landmarks or vice versa.
        """
        if (not(type(size)) is list) or (len(size) != 3): raise Exception("size must be a list of length 3.")
        lmkList = []
        labelList = self.GetLabels()
        s = mat(size)
        for i in range(self.GetNumberOfLandmarks()):
            x0 = mat(self.landmarkList[i][1:])
            x1 = s - 1 - x0
            lmk  = [labelList[i]] + x1.tolist()[0]
            lmkList.append(lmk)
        return landmarks(lmkList, self.spacing)

    """
    def Resample(self, inSpacing, outSpacing):
        outLandmarkList = []
        for inLandmark in self.landmarkList:
            inLabel = inLandmark[0]
            inIndex = mat(inLandmark[1:])
            x = multiply(inIndex, inSpacing)
            outIndex = x / outSpacing
            outLandmark = [inLabel] + outIndex.tolist()[0]
            outLandmarkList.append(outLandmark)

        return landmarks(outLandmarkList)
    """

    def Crop(self, size):
        if (not(type(size)) is list) or (len(size) != 3): raise Exception("size must be a list of length 3.")
        outLandmarkList = []
        for inLandmark in self.landmarkList:
            inLabel = inLandmark[0]
            inCoordinate = mat(inLandmark[1:])
            outCoordinate = inCoordinate - mat(size)
            outLandmark = [inLabel] + outCoordinate.tolist()[0]
            outLandmarkList.append(outLandmark)

        return landmarks(outLandmarkList, self.spacing)

class surface(landmarks):
    def __init__(self, inputPoints=[[0,0,0]], spacing=[1.0,1.0,1.0], inputConnections=None):
        self.numSurfaces = 1
        self.numPoints = 1
        self.numPolygons = 0
        self.numConnections = 0
        self.connectionList = []
        if type(inputPoints) is list:
            inputLandmarks = []
            for (i,point) in enumerate(inputPoints):
                landmark = [str(i)] + point
                inputLandmarks+=[landmark]
            landmarks.__init__(self, inputLandmarks)
        elif type(inputPoints) is str:
            self.Read(inputPoints, spacing)

    def Write(self, path):
        """
        Writes surface to given path using BYU format
        """
        # Write metadata
        surfaceFile = open(path,"w")
        print("{0} {1} {2} {3}".format(self.numSurfaces, self.GetNumberOfPoints(), self.numPolygons, self.GetNumberOfConnections()), file=surfaceFile)

        for i in range(self.numSurfaces):
            print("{0} {1}".format(self.surfaceConnectionsStart[i], self.surfaceConnectionsEnd[i]),file=surfaceFile)

        # Write points
        pointList = self.GetPoints()
        for point in pointList:
            print(" ".join(map(str,point)),file=surfaceFile)

        # Write connections
        for connection in self.connectionList[:]:
            print(" ".join(map(str,connection[:-1]))+" -"+str(connection[-1]), file=surfaceFile)

        surfaceFile.close()

    def Read(self, path, spacing=[1.0,1.0,1.0]):
        """
        Reads surface from given path
        It accepts surfaces in the following format
        1. BYU format which has extension .byu

           <numSurfaces> <numPoints> <numPolygons> <numConnections>
           1 <numPointsInSurface1>
           2 <numPointsInSurface2>
           ...
           <numSurfaces> <numPointsInSurface<numSurfaces>>
           x1 y1 z1
           x2 y2 z2
           ...
           x<numPoints> y<numPoints> z<numPoints>
        """
        dimension = 3
        self.spacing = spacing
        surfaceFile = open(path, "r")
        lineList = []
        for line in surfaceFile.readlines():
            if line.strip() != "": lineList.append(line.strip())
        surfaceFile.close()

        # Read Metadata
        line1 = lineList.pop(0).split()
        [self.numSurfaces, self.numPoints, self.numPolygons, self.numConnections] = map(int, line1[0:4])

        self.surfaceConnectionsStart = []
        self.surfaceConnectionsEnd = []
        for i in range(self.numSurfaces):
            [start, end] = lineList.pop(0).split()
            self.surfaceConnectionsStart.append(int(start))
            self.surfaceConnectionsEnd.append(int(end))

        # Read Points
        pointData = []
        while len(pointData) < self.numPoints*dimension:
            line = lineList.pop(0)
            pointData+=map(float,line.split())

        landmarkList = []
        for i in range(self.numPoints):
            point = pointData[i*dimension:(i+1)*dimension]
            label = str(i)
            landmark = [label] + point
            landmarkList.append(landmark)

        self.SetLandmarks(landmarkList)

        # Read Connections
        connectionData = []
        for line in lineList: connectionData += map(int,line.split())

        self.connectionList = []
        connection = []
        for value in connectionData:
            if value > 0:
                connection.append(value)
            else:
                connection.append(-value)
                self.connectionList.append(connection)
                connection=[]

    def GetNumberOfConnections(self):
        return len(self.connectionList)

    def Affine(self, affine):
        print(type(self))
        ### TODO generalize to other dimensions
        if (not(type(affine)) is list) or (len(affine) != 12): raise Exception("affine must be a list of length 12.")
        pointList = []
        for i in range(self.GetNumberOfPoints()):
            x0 = mat(self.landmarkList[i][1:]).reshape(3,1)
            A = mat(affine[:9]).reshape(3,3)
            b = mat(affine[9:]).reshape(3,1)
            #x1 = A.I*(x0 - b)
            x1 = A*x0 + b
            point = x1.flatten().tolist()[0]
            pointList.append(point)
        return surface(pointList, spacing=self.spacing)


#!/usr/bin/python
"""
This script is for processing the ADV data files in the ADV/ directory.
"""
import dolfyn.adv.api as avm
from dolfyn.tools import within
from dolfyn.data.time import num2date
import numpy as np
from os.path import isfile
from main import FILEINFO

# # The file names:
# FNAMES = {ky: val.basename for ky, val in FILEINFO.iteritems()}

# Some variables for calculating dissipation rate (\epsilon)
eps_freqs = np.array([[.3, 1],
                      [.3, 1],
                      [.3, 3], ])
spec_noise = [1.5e-4,
              1.5e-4,
              1.5e-5, ]
pii = 2 * np.pi

mc = avm.motion.CorrectMotion()

def run(finfo=FILEINFO.values(), readvec=None, savecsv=False):
    """
    Process the ADV data.

    Parameters
    ----------
    fnames : iterable
         A list of data file names that you want to process (default: all of
         them).
    readvec : {True, None, False}
         Whether to read the raw ``.vec`` file, or load the ``.h5``
         file. Default: read .h5, if it is available.
    savecsv : bool
         Save the ``_average5min.csv`` files?
    """
    for finf in finfo:
        print("File: {}".format(finf.fname))
        if readvec is True or \
           readvec is None and not isfile(finf.abs_fname + '.h5'):
            dr = _read_raw(finf)
        else:
            dr = avm.load(finf.abs_fname + '.h5')

        drm = correct_motion(dr, finf)

        print('  Saving matlab file...')
        drm.add_data('datenum', drm.mpltime + 366, 'main')
        drm.save_mat(finf.abs_fname + '_earth.mat', groups=['orient', 'main'])
        drm.pop_data('datenum')

        bdat = average(drm)

        print("  Saving binned data to hdf5...")
        bdat.save(finf.abs_fname + '_earth_b5m.h5')

        if savecsv:
            _save_csv(bdat, finf)

        print("  Rotating to Principal frame...")
        avm.rotate.earth2principal(drm)
        print("  Binning and saving...")
        bdat2 = average(drm)

        bdat2.save(finf.abs_fname + '_pax_b5m.h5')

        print("Done.")

def _read_raw(finf):
    # Read the raw vector file
    if not isfile(finf.abs_source_fname):
        print("File not found.")
        if finf in FILEINFO.values():
            print("... Try running the main.pull function?")
    dr = avm.read_nortek(finf.abs_source_fname)

    dr.noise[0] = 0
    dr.noise[1] = 0
    dr.noise[2] = 0

    # Crop the data when the instrument was on the seafloor
    dr = dr.subset(within(dr.mpltime, dr.props['time_range']))

    ##########
    print('  Cleaning the data...')
    dr.u[~within(dr.u, [-2.5, 0.5])] = np.NaN
    avm.clean.fillpoly(dr.u, 3, 12)
    dr.v[~within(dr.v, [-1, 1])] = np.NaN
    avm.clean.fillpoly(dr.v, 3, 12)
    dr.v[~within(dr.w, [-1, 1])] = np.NaN
    avm.clean.fillpoly(dr.w, 3, 12)
    avm.clean.GN2002(dr.u)
    avm.clean.GN2002(dr.v)
    avm.clean.GN2002(dr.w)
    for nm, d in dr.iter():
        if isinstance(d, np.ndarray) and \
           d.dtype == np.float64 and nm not in ['mpltime']:
            dr[nm] = d.astype(np.float32)
    if dr.has_imu:
        (dr.pitch,
         dr.roll,
         dr.heading) = avm.rotate.orient2euler(dr.orientmat)
    print('  Saving...')
    dr.save(finf.abs_fname + '.h5',
            units={
                'vel': 'm/s', 'velrot': 'm/s',
                'velacc': 'm/s', 'AngRt': 'rad/s',
                'Accel': 'm/s^2', 'AccelStable': 'm/s^2',
                'pitch': 'deg', 'roll': 'deg', 'heading': 'deg true',
                'mpltime': 'MatPlotLib Time format',
                'time': 'ISO8601 time strings.', },
            description={
                'vel': 'Velocity array 0:True East, 1: True North, 2: Up',
                'velrot': 'The rotation-rate velocity.',
                'velacc': 'The tranlational (acceleration) velocity.',
                'AccelStable': 'The low-frequency acceleration '
                'that is ignored in calculating velacc.',
                'pitch': 'The pitch of the ADV body',
                'roll': 'The roll of the ADV body',
                'heading': 'The heading (True) of the ADV body',
                'orientmat': "The orientation matrix of the"
                " ADV body in the Earth's reference frame", })
    return dr

def _save_csv(bdat, finf):

        ti = bdat.sigma_Uh / bdat.U_mag
        ti[bdat.U_mag < 0.7] = np.NaN

        print("  Saving average csv file...")
        np.savetxt(finf.abs_fname + '_Average5min.csv',
                   np.vstack((num2date(bdat.mpltime), bdat.u,
                              bdat.v, bdat.w, ti)).T,
                   fmt=['%s'] + ['%0.3f'] * 4,
                   header='Date+Time (US/Pacific), u (true east m/s), '
                          'v (true north m/s), w (up m/s), '
                          'Turbulence Intensity',
                   delimiter=', ')

def correct_motion(dr, finf):

    print('  Motion correcting...')
    drm = dr.copy()
    mc(drm)
    (drm.pitch[:],
     drm.roll[:],
     drm.heading[:]) = avm.rotate.orient2euler(drm.orientmat)
    print('  Saving...')
    drm.save(finf.abs_fname + '_earth.h5',
             units={
                 'vel': 'm/s', 'velrot': 'm/s',
                 'velacc': 'm/s', 'AngRt': 'rad/s',
                 'Accel': 'm/s^2', 'AccelStable': 'm/s^2',
                 'pitch': 'deg', 'roll': 'deg', 'heading': 'deg true',
                 'mpltime': 'MatPlotLib Time format',
                 'time': 'ISO8601 time strings.', },
             description={
                 'vel': 'Velocity array 0:True East, 1: True North, 2: Up',
                 'velrot': 'The rotation-rate velocity.',
                 'velacc': 'The tranlational (acceleration) velocity.',
                 'AccelStable': 'The low-frequency acceleration that'
                 ' is ignored in calculating velacc.',
                 'pitch': 'The pitch of the ADV body',
                 'roll': 'The roll of the ADV body',
                 'heading': 'The heading (True) of the ADV body',
                 'orientmat': "The orientation matrix of the ADV "
                 "body in the Earth's reference frame", })
    return drm

def average(dat):
    print("  Averaging...")
    bnr = avm.TurbBinner(n_bin=5 * 60 * dat.fs, fs=dat.fs)

    # This is just a shortcut
    velmot = dat.velacc + dat.velrot

    bdat = bnr(dat)

    # Calculate spectra ####
    bdat.add_data('Spec_velrot',
                  bnr.calc_vel_psd(dat.velrot, ),
                  'spec')
    bdat.add_data('Spec_velacc',
                  bnr.calc_vel_psd(dat.velacc, ),
                  'spec')
    bdat.add_data('Spec_velmot',
                  bnr.calc_vel_psd(velmot, ),
                  'spec')
    bdat.add_data('Spec_velraw',
                  bnr.calc_vel_psd(dat.velraw, ),
                  'spec')

    # Calculate cross-spectra ('point cross-spectra') ####
    bdat.props['Cspec_comp'] = ['uv', 'uw', 'vw']
    bdat.add_data('Cspec_vel',
                  bnr.calc_vel_cpsd(dat.vel).astype(np.complex64),
                  'spec')
    bdat.add_data('Cspec_velmot',
                  bnr.calc_vel_cpsd(velmot).astype(np.complex64),
                  'spec')
    bdat.add_data('Cspec_velraw',
                  bnr.calc_vel_cpsd(dat.velraw).astype(np.complex64),
                  'spec')

    # Calculate triple products ####
    # setup
    bdat.props['tripprod_comp'] = [['uuu', 'uuv', 'uuw', ],
                                   ['vvu', 'vvv', 'vvw', ],
                                   ['wwu', 'wwv', 'www', ]]
    bdat.add_data('tripprod',
                  np.empty((3, 3, len(bdat.u)),
                           dtype=np.float32),
                  'turb')
    # Calculate
    turb = bnr.demean(dat.vel)
    for i0 in range(3):
        for i1 in range(3):
            bdat.tripprod[i0, i1] = (turb[i0] ** 2 * turb[i1]).mean(-1)

    # Calculate the dissipation rate ####
    epstmp = np.zeros_like(bdat.u)
    Ntmp = 0
    for idx, frq_rng in enumerate(eps_freqs):
        if frq_rng is None:
            continue
        om_rng = frq_rng * pii
        N = ((om_rng[0] < bdat.omega) & (bdat.omega < om_rng[1])).sum()
        sptmp = bdat.Spec[idx] - spec_noise[idx] / pii
        sptmp[sptmp < 0] = 0
        tmp = bnr.calc_epsilon_LT83(sptmp,
                                    bdat.omega,
                                    np.abs(bdat.U),
                                    om_rng)
        epstmp += tmp * N
        Ntmp += N
    epstmp /= Ntmp
    # epstmp[np.abs(dat.U) < 0.2] = np.NaN
    bdat.add_data('epsilon', epstmp, 'main')
    return bdat

if __name__ == '__main__':

    import argparse

    parser = argparse.ArgumentParser(
        description="Process ADV data files.")
    parser.add_argument(
        'fnames', nargs='*',
        help='The base file names (<basename>, i.e., without '
        'file extensions) of the files to process. If no files '
        'are specified, all files specified in the scrip will be processed.')
    parser.add_argument(
        '--readvec',
        help="Force reading of the binary vector file. By default, "
        "the script will only read `<basename>.vec` files if there "
        "is no <basename>.h5",
        action='store_true')
    parser.add_argument(
        '--savecsv',
        help="Save simplified CSV files during processing?",
        action='store_true')
    args = parser.parse_args()
    if not args.readvec:
        args.readvec = None
    if len(args.fnames) == 0:
        args.fnames = FNAMES.values()

    run(args.fnames, readvec=args.readvec, savecsv=args.savecsv)