import os
import numpy as np
import subprocess
from ..mbase import which
from ..utils.cvfdutil import centroid_of_polygon
from ..plot.plotutil import plot_cvfd

class Triangle(object):
    """
    Class to work with the triangle program to unstructured triangular grids.
    Information on the triangle program can be found at
    https://www.cs.cmu.edu/~quake/triangle.html

    Parameters
    ----------
    model_ws : str
        workspace location for creating triangle files (default is '.')
    exe_name : str
        path and name of the triangle program. (default is triange, which
        means that the triangle program must be in your path)
    maximum_area : float
        the maximum area for any triangle.  The default value is None, which
        means that the user must specify maximum areas for each region.
    angle : float
        Triangle will continue to add vertices until no angle is less than
        this specified value.  (default is 20 degrees)
    additional_args : list
        list of additional command line switches to pass to triangle

    Returns
    -------
    None

    """
    def __init__(self, model_ws='.', exe_name='triangle', maximum_area=None,
                 angle=20., additional_args=None):
        self.model_ws = model_ws
        exe_name = which(exe_name)
        if exe_name is None:
            raise Exception('Cannot find gridgen binary executable')
        self.exe_name = os.path.abspath(exe_name)
        self.angle = angle
        self.maximum_area = maximum_area
        self.additional_args = additional_args
        self._initialize_vars()
        return

    def add_polygon(self, polygon):
        """
        Add a polygon

        Parameters
        ----------
        polygon : list
            polygon is a list of (x, y) points

        Returns
        -------
        None

        """
        self._polygons.append(polygon)
        return

    def add_hole(self, hole):
        """
        Add a point that will turn enclosing polygon into a hole

        Parameters
        ----------
        hole : tuple
            (x, y)

        Returns
        -------
        None

        """
        self._holes.append(hole)
        return

    def add_region(self, point, attribute=0, maximum_area=None):
        """
        Add a point that will become a region with a maximum area, if
        specified.

        Parameters
        ----------
        point : tuple
            (x, y)

        attribute : integer or float
            integer value assigned to output elements

        maximum_area : float
            maximum area of elements in region

        Returns
        -------
        None

        """
        self._regions.append([point, attribute, maximum_area])
        return

    def build(self, verbose=False):
        """
        Build the triangular mesh

        Parameters
        ----------
        verbose : bool
            If true, print the results of the triangle command to the terminal
            (default is False)

        Returns
        -------
        None

        """

        # provide some protection by removing existing files
        self.clean()

        # write the active domain to a file
        fname = os.path.join(self.model_ws, self.file_prefix + '.0.node')
        self._write_nodefile(fname)

        # poly file
        fname = os.path.join(self.model_ws, self.file_prefix + '.0.poly')
        self._write_polyfile(fname)

        # Construct the triangle command
        cmds = [self.exe_name]
        if self.maximum_area is not None:
            cmds.append('-a{}'.format(self.maximum_area))
        else:
            cmds.append('-a')
        if self.angle is not None:
            cmds.append('-q{}'.format(self.angle))
        if self.additional_args is not None:
            cmds += self.additional_args
        cmds.append('-A') # assign attributes
        cmds.append('-p') # triangulate .poly file
        cmds.append('-V') # verbose
        cmds.append('-D') # delaunay triangles for finite volume
        cmds.append('-e') # edge file
        cmds.append('-n') # neighbor file
        cmds.append(self.file_prefix + '.0') # output file name

        # run Triangle
        buff = subprocess.check_output(cmds, cwd=self.model_ws)
        buff = buff.decode()
        if verbose:
            print(buff)

        # load the results
        self._load_results()
        self.ncpl = self.ele.shape[0]
        self.nvert = self.node.shape[0]

        # create verts and iverts
        self.verts = self.node[['x', 'y']]
        self.verts = np.array(self.verts.tolist(), np.float)
        self.iverts = []
        for row in self.ele:
            self.iverts.append([row[1], row[2], row[3]])

        return

    def plot(self, ax=None, layer=0, edgecolor='k', facecolor='none',
             cmap='Dark2', a=None, masked_values=None, **kwargs):
        """
        Plot the grid.  This method will plot the grid using the shapefile
        that was created as part of the build method.

        Note that the layer option is not working yet.

        Parameters
        ----------
        ax : matplotlib.pyplot axis
            The plot axis.  If not provided it, plt.gca() will be used.
            If there is not a current axis then a new one will be created.
        layer : int
            Layer number to plot
        cmap : string
            Name of colormap to use for polygon shading (default is 'Dark2')
        edgecolor : string
            Color name.  (Default is 'scaled' to scale the edge colors.)
        facecolor : string
            Color name.  (Default is 'scaled' to scale the face colors.)
        a : numpy.ndarray
            Array to plot.
        masked_values : iterable of floats, ints
            Values to mask.
        kwargs : dictionary
            Keyword arguments that are passed to
            PatchCollection.set(``**kwargs``).  Some common kwargs would be
            'linewidths', 'linestyles', 'alpha', etc.

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt

        if ax is None:
            ax = plt.gca()

        pc = plot_cvfd(self.verts, self.iverts, ax=ax, edgecolor=edgecolor,
                       facecolor=facecolor, cmap=cmap, a=a,
                       masked_values=masked_values, **kwargs)
        ax.autoscale()
        return pc

    def get_boundary_marker_array(self):
        """
        Get an integer array that has boundary markers

        Returns
        -------
        iedge : ndarray
            integer array of size ncpl containing a boundary ids.  The array
            contains zeros for cells that do not touch a boundary.  The
            boundary ids are the segment numbers for each segment in each
            polygon that is added with the add_polygon method.

        """
        iedge = np.zeros((self.ncpl), dtype=np.int)
        boundary_markers = np.unique(self.edge['boundary_marker'])
        for ibm in boundary_markers:
            icells = self.get_edge_cells(ibm)
            iedge[icells] = ibm
        return iedge

    def plot_boundary(self, ibm, ax=None, **kwargs):
        """
        Plot a line and vertices for the specified boundary marker

        Parameters
        ----------
        ibm : integer
            plot the boundary for this boundary marker

        ax : matplotlib.pyplot.Axes
           axis to add the plot to.  (default is plt.gca())

        kwargs : dictionary
            dictionary of arguments to pass to ax.plot()

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        idx = np.where(self.edge['boundary_marker'] == ibm)[0]
        for i in idx:
            iv1 = self.edge['endpoint1'][i]
            iv2 = self.edge['endpoint2'][i]
            x1 = self.node['x'][iv1]
            x2 = self.node['x'][iv2]
            y1 = self.node['y'][iv1]
            y2 = self.node['y'][iv2]
            ax.plot([x1, x2], [y1, y2], **kwargs)
        return

    def plot_vertices(self, ax=None, **kwargs):
        """
        Plot the mesh vertices

        Parameters
        ----------
        ax : matplotlib.pyplot.Axes
           axis to add the plot to.  (default is plt.gca())

        kwargs : dictionary
            dictionary of arguments to pass to ax.plot()

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        ax.plot(self.node['x'], self.node['y'], lw=0, **kwargs)
        return

    def label_vertices(self, ax=None, onebased=True, **kwargs):
        """
        Label the mesh vertices with their vertex numbers

        Parameters
        ----------
        ax : matplotlib.pyplot.Axes
           axis to add the plot to.  (default is plt.gca())

        onebased : bool
            Make the labels one-based if True so that they correspond to
            what would be written to MODFLOW.

        kwargs : dictionary
            dictionary of arguments to pass to ax.text()

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        for i in range(self.verts.shape[0]):
            x = self.verts[i, 0]
            y = self.verts[i, 1]
            s = i
            if onebased:
                s += 1
            s = '{}'.format(s)
            ax.text(x, y, s, **kwargs)
        return

    def plot_centroids(self, ax=None, **kwargs):
        """
        Plot the cell centroids

        Parameters
        ----------
        ax : matplotlib.pyplot.Axes
           axis to add the plot to.  (default is plt.gca())

        kwargs : dictionary
            dictionary of arguments to pass to ax.plot()

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        xcyc = self.get_xcyc()
        ax.plot(xcyc[:, 0], xcyc[:, 1], lw=0, **kwargs)
        return

    def label_cells(self, ax=None, onebased=True, **kwargs):
        """
        Label the cells with their cell numbers

        Parameters
        ----------
        ax : matplotlib.pyplot.Axes
           axis to add the plot to.  (default is plt.gca())

        onebased : bool
            Make the labels one-based if True so that they correspond to
            what would be written to MODFLOW.

        kwargs : dictionary
            dictionary of arguments to pass to ax.text()

        Returns
        -------
        None

        """
        import matplotlib.pyplot as plt
        if ax is None:
            ax = plt.gca()
        xcyc = self.get_xcyc()
        for i in range(xcyc.shape[0]):
            x = xcyc[i, 0]
            y = xcyc[i, 1]
            s = i
            if onebased:
                s += 1
            s = '{}'.format(s)
            ax.text(x, y, s, **kwargs)
        return

    def get_xcyc(self):
        """
        Get a 2-dimensional array of x and y cell center coordinates.

        Returns
        -------
        xcyc : ndarray
            column 0 contains the x coordinates and column 1 contains the
            y coordinates

        """
        ncpl = len(self.iverts)
        xcyc = np.empty((ncpl, 2), dtype=np.float)
        for i, icell2d in enumerate(self.iverts):
            points = []
            for iv in icell2d:
                x = self.verts[iv, 0]
                y = self.verts[iv, 1]
                points.append((x, y))
            xc, yc = centroid_of_polygon(points)
            xcyc[i, 0] = xc
            xcyc[i, 1] = yc
        return xcyc

    def get_cell2d(self):
        """
        Get a list of the information needed for the MODFLOW DISV Package.

        Returns
        -------
        cell2d : list (of lists)
            innermost list contains cell number, x, y, number of vertices, and
            then the vertex numbers comprising the cell.

        """
        cell2d = []
        xcyc = self.get_xcyc()
        for i, icell2d in enumerate(self.iverts):
            ic2dr = icell2d[::-1]
            cell2d.append([i, xcyc[i, 0], xcyc[i, 1], len(icell2d)] + ic2dr)
        return cell2d

    def get_vertices(self):
        """
        Get a list of vertices in the form needed for the MODFLOW DISV Package.

        Returns
        -------
        vertices : list (of lists)
            innermost list contains vertex number, x, and y

        """
        vertices = []
        for i, row in enumerate(self.verts):
            vertices.append([i, row[0], row[1]])
        return vertices

    def get_edge_cells(self, ibm):
        """
        Get a list of cell numbers that correspond to the specified boundary
        marker.

        Parameters
        ----------
        ibm : integer
            boundary marker value

        Returns
        -------
        cell_list : list
            list of zero-based cell numbers

        """
        # Create the edge dictionary if it doesn't exist
        if self.edgedict is None:
            edgedict = {}
            for ie, iv1, iv2, iseg in self.edge:
                if iseg != 0:
                    edgedict[(iv1, iv2)] = iseg
                    edgedict[(iv2, iv1)] = iseg

        # Create a list of cells for boundary marker ibm
        cell_list = []
        for n, ivlist in enumerate(self.iverts):
            itmp = ivlist + [ivlist[0]]
            for i in range(len(ivlist)):
                ie = (itmp[i], itmp[i + 1])
                if ie in edgedict:
                    if edgedict[ie] == ibm:
                        cell_list.append(n)

        return cell_list

    def get_attribute_array(self):
        """
        Return an array containing the attribute value for each cell.  These
        are the attribute values that are passed into the add_region() method.

        Returns
        -------
        attribute_array : ndarray

        """
        return self.ele['attribute']

    def clean(self):
        """
        Remove the input and output files created by this class and by the
        Triangle program

        Returns
        -------
        None

        """
        # remove input files
        for ext in ['poly', 'node']:
            fname = os.path.join(self.model_ws, self.file_prefix + '0.' + ext)
            if os.path.isfile(fname):
                os.remove(fname)
                if os.path.isfile(fname):
                    print('Could not remove: {}'.format(fname))
        # remove output files
        for ext in ['poly', 'ele', 'node', 'neigh', 'edge']:
            fname = os.path.join(self.model_ws, self.file_prefix + '1.' + ext)
            if os.path.isfile(fname):
                os.remove(fname)
                if os.path.isfile(fname):
                    print('Could not remove: {}'.format(fname))
        return

    def _initialize_vars(self):
        self.file_prefix = '_triangle'
        self.ncpl = 0
        self.nvert = 0
        self._active_domain = None
        self._polygons = []
        self._holes = []
        self._regions = []
        self.verts = None
        self.iverts = None
        self.edgedict = None
        return

    def _load_results(self):

        # node file
        ext = 'node'
        dt = [('ivert', int), ('x', float), ('y', float)]
        fname = os.path.join(self.model_ws, self.file_prefix + '.1.' + ext)
        setattr(self, ext, None)
        if os.path.isfile(fname):
            f = open(fname, 'r')
            line = f.readline()
            f.close()
            ll = line.strip().split()
            nvert = int(ll[0])
            ndim = int(ll[1])
            assert ndim == 2, 'Dimensions in node file is not 2'
            iattribute = int(ll[2])
            if iattribute == 1:
                dt.append(('attribute', int))
            ibm = int(ll[3])
            if ibm == 1:
                dt.append(('boundary_marker', int))
            a = np.loadtxt(fname, skiprows=1, comments='#', dtype=dt)
            assert a.shape[0] == nvert
            setattr(self, ext, a)

        # ele file
        ext = 'ele'
        dt = [('icell', int), ('iv1', int), ('iv2', int), ('iv3', int)]
        fname = os.path.join(self.model_ws, self.file_prefix + '.1.' + ext)
        setattr(self, ext, None)
        if os.path.isfile(fname):
            f = open(fname, 'r')
            line = f.readline()
            f.close()
            ll = line.strip().split()
            ncells = int(ll[0])
            npt = int(ll[1])
            assert npt == 3, 'Nodes per triangle in ele file is not 3'
            iattribute = int(ll[2])
            if iattribute == 1:
                dt.append(('attribute', int))
            a = np.loadtxt(fname, skiprows=1, comments='#', dtype=dt)
            assert a.shape[0] == ncells
            setattr(self, ext, a)

        # edge file
        ext = 'edge'
        dt = [('iedge', int), ('endpoint1', int), ('endpoint2', int)]
        fname = os.path.join(self.model_ws, self.file_prefix + '.1.' + ext)
        setattr(self, ext, None)
        if os.path.isfile(fname):
            f = open(fname, 'r')
            line = f.readline()
            f.close()
            ll = line.strip().split()
            nedges = int(ll[0])
            ibm = int(ll[1])
            if ibm == 1:
                dt.append(('boundary_marker', int))
            a = np.loadtxt(fname, skiprows=1, comments='#', dtype=dt)
            assert a.shape[0] == nedges
            setattr(self, ext, a)

        # neighbor file
        ext = 'neigh'
        dt = [('icell', int), ('neighbor1', int), ('neighbor2', int),
              ('neighbor3', int)]
        fname = os.path.join(self.model_ws, self.file_prefix + '.1.' + ext)
        setattr(self, ext, None)
        if os.path.isfile(fname):
            f = open(fname, 'r')
            line = f.readline()
            f.close()
            ll = line.strip().split()
            ncells = int(ll[0])
            nnpt = int(ll[1])
            assert nnpt == 3, 'Neighbors per triangle in neigh file is not 3'
            a = np.loadtxt(fname, skiprows=1, comments='#', dtype=dt)
            assert a.shape[0] == ncells
            setattr(self, ext, a)

        return

    def _write_nodefile(self, fname):
        f = open(fname, 'w')
        nvert = 0
        for p in self._polygons:
            nvert += len(p)
        s = '{} {} {} {}\n'.format(nvert, 2, 0, 0)
        f.write(s)
        ip = 0
        for p in self._polygons:
            for i, vertex in enumerate(p):
                s = '{} {} {}\n'.format(ip, vertex[0], vertex[1])
                f.write(s)
                ip += 1
        f.close()

    def _write_polyfile(self, fname):
        f = open(fname, 'w')

        # vertices, write zero to indicate read from node file
        s = '{} {} {} {}\n'.format(0, 0, 0, 0)
        f.write(s)

        # segments
        nseg = 0
        for p in self._polygons:
            nseg += len(p)
        bm = 1
        s = '{} {}\n'.format(nseg, bm)
        f.write(s)

        iseg = 0
        ipstart = 0
        for p in self._polygons:
            nseg = len(p)
            for i in range(nseg):
                ep1 = i
                ep2 = i + 1
                if ep2 > nseg - 1:
                    ep2 = 0
                ep1 += ipstart
                ep2 += ipstart
                s = '{} {} {} {}\n'.format(iseg, ep1, ep2, iseg + 1)
                f.write(s)
                iseg += 1
            ipstart += len(p)

        # holes
        nholes = len(self._holes)
        s = '{}\n'.format(nholes)
        f.write(s)
        for i, hole in enumerate(self._holes):
            s = '{} {} {}\n'.format(i, hole[0], hole[1])
            f.write(s)

        # regions
        nregions = len(self._regions)
        s = '{}\n'.format(nregions)
        f.write(s)
        for i, region in enumerate(self._regions):
            pt = region[0]
            attribute = region[1]
            maxarea = region[2]
            if maxarea is None:
                maxarea = -1.
            s = '{} {} {} {} {}\n'.format(i, pt[0], pt[1], attribute, maxarea)
            f.write(s)

        f.close()
        return


import os
import sys
import copy
import numpy
import weakref

from astropy.coordinates import EarthLocation
import astropy.units as u

__all__ = ['OVRO_CONFIG_FILENAME', 'Station', 'Antenna', 'parse_config', 'ovro']

OVRO_CONFIG_FILENAME = os.path.join(os.path.dirname(__file__), 'ovro.txt')

def _smart_int(s, positive_only=False):
    i = 0
    v = None
    while i < len(s):
        try:
            if positive_only and s[i] == '-':
                raise ValueError
            v = int(s[i:], 10)
            break
        except ValueError:
            pass
        i += 1
    if v is None:
        raise ValueError("Cannot convert '%s' to int" % s)
    return v

class Station(object):
    """
    Class to represent the OVRO-LWA station and its antennas.
    """

    def __init__(self, name, lat, lon, elev, antennas=None):
        self.name = name
        self.lat = lat
        self.lon = lon
        self.elev = elev

        self.antennas = []
        if antennas is not None:
            self.antennas = antennas

    @classmethod
    def from_line(cls, line):
        """
        Create a new Station instance from a line in an antenna positions file.
        """

        name, style, lat, lon, elev, x, y, active = line.split(None, 7)
        lat = float(lat) * numpy.pi/180
        lon = float(lon) * numpy.pi/180
        elev = 1182.89  ## Mean of the first 256 antennas
        return cls(name, lat, lon, elev)

    def append(self, ant):
        """
        Add an antenna to the array.
        """

        if not isinstance(ant, Antenna):
            raise TypeError("Expected an antenna")
        ant.parent = weakref.proxy(self)
        self.antennas.append(ant)

    def select_subset(self, ids):
        """
        Given a list of antenna IDs (either as integer index or name), return a
        new Station instance that only contains those antennas.
        """

        subset = Station(self.name+"-fast", self.lat*1.0, self.lon*1.0, self.elev*1.0)

        all_ids = [ant.id for ant in self.antennas]
        for id in ids:
            if isinstance(id, int):
                subset.append(copy.deepcopy(self.antennas[id]))
            else:
                subset.append(copy.deepcopy(self.antennas[all_ids.index(id)]))
        return subset

    @property
    def ecef(self):
        """
        Return the Earth centered, Earth fixed location of the array in meters.
        """

        e =  EarthLocation(lat=self.lat*u.rad, lon=self.lon*u.rad, height=self.elev*u.m)
        return (e.x.to_value(u.m), e.y.to_value(u.m), e.z.to_value(u.m))

    @property
    def topo_rot_matrix(self):
        """
        Return the rotation matrix that takes a difference in an Earth centered,
        Earth fixed location relative to the Station and rotates it into a
        topocentric frame that is south-east-zenith.
        """

        r = numpy.array([[ numpy.sin(self.lat)*numpy.cos(self.lon), numpy.sin(self.lat)*numpy.sin(self.lon), -numpy.cos(self.lat)],
                         [-numpy.sin(self.lon),                     numpy.cos(self.lon),                      0                  ],
                         [ numpy.cos(self.lat)*numpy.cos(self.lon), numpy.cos(self.lat)*numpy.sin(self.lon),  numpy.sin(self.lat)]])
        return r

    @property
    def casa_position(self):
        """
        Return a four-element tuple of (CASA position reference, CASA position 1,
        CASA position 2, CASA position 3, CASA position 4) that is suitable for
        use with casacore.measures.measures.position.
        """

        x, y, z = self.ecef
        return 'ITRF', '%fm' % x, '%fm' % y, '%fm' % z

class Antenna(object):
    """
    Class to represent an antenna in the OVRO-LWA.
    """

    def __init__(self, id, lat, lon, elev):
        if isinstance(id, str):
            id = _smart_int(id, positive_only=True)
        self.id = id
        self.lat = lat
        self.lon = lon
        self.elev = elev
        self.parent = None

    @classmethod
    def from_line(cls, line):
        """
        Create a new Antenna instance from a line in an antenna positions file.
        """

        name, style, lat, lon, elev, x, y, active = line.split(None, 7)
        lat = float(lat) * numpy.pi/180
        lon = float(lon) * numpy.pi/180
        try:
            elev = float(elev)
        except ValueError:
            elev = 1182.89  ## Mean of the first 256 antennas
        return cls(name, lat, lon, elev)

    @property
    def ecef(self):
        """
        Return the Earth centered, Earth fixed location of the antenna in meters.
        """

        e = EarthLocation(lat=self.lat*u.rad, lon=self.lon*u.rad, height=self.elev*u.m)
        return (e.x.to_value(u.m), e.y.to_value(u.m), e.z.to_value(u.m))

    @property
    def enz(self):
        """
        Return the topocentric east-north-zenith coordinates for the antenna
        relative to the center of its associated Station in meters.
        """

        if self.parent is None:
            raise RuntimeError("Cannot find east-north-zenith without an associated Station")

        ecefFrom = numpy.array(self.parent.ecef)
        ecefTo = numpy.array(self.ecef)

        rho = ecefTo - ecefFrom
        rot = self.parent.topo_rot_matrix
        sez = numpy.dot(rot, rho)

        # Convert from south, east, zenith to east, north, zenith
        enz = 1.0*sez[[1,0,2]]
        enz[1] *= -1.0
        return enz

def parse_config(filename):
    """
    Given an OVRO-LWA configuration file, parse it and return a Station instance.
    """

    with open(filename, 'r') as fh:
        for line in fh:
            if len(line) < 3:
                continue
            elif line[0] == '#':
                continue

            if line.startswith('LWA-000'):
                line = line.replace('LWA-000', 'OVRO-LWA')
                station = Station.from_line(line)
            elif line.find('NO') > -2:
                ant = Antenna.from_line(line)
                station.append(ant)

    return station

# A ready-made Station instance, filled with Antennas
ovro = parse_config(OVRO_CONFIG_FILENAME)

# Use OVRO_MMA as the telescope name until CASA knows about OVRO-LWA
ovro.name = 'OVRO_MMA'

# Change the order to match what's going on in Phase I
## The current list as of 2022 Jan 5
interim = ['LWA-266', 'LWA-259', 'LWA-268', 'LWA-267', 'LWA-271', 'LWA-269',
           'LWA-276', 'LWA-273', 'LWA-278', 'LWA-277', 'LWA-282', 'LWA-281',
           'LWA-307', 'LWA-285', 'LWA-309', 'LWA-308', 'LWA-311', 'LWA-310',
           'LWA-313', 'LWA-312', 'LWA-321', 'LWA-314', 'LWA-330', 'LWA-327',
           'LWA-338', 'LWA-332', 'LWA-340', 'LWA-339', 'LWA-352', 'LWA-341',
           'LWA-362', 'LWA-353', 'LWA-257', 'LWA-255', 'LWA-260', 'LWA-258',
           'LWA-265', 'LWA-263', 'LWA-272', 'LWA-270', 'LWA-283', 'LWA-280',
           'LWA-288', 'LWA-284', 'LWA-292', 'LWA-291', 'LWA-296', 'LWA-295',
           'LWA-301', 'LWA-298', 'LWA-305', 'LWA-303', 'LWA-317', 'LWA-306',
           'LWA-320', 'LWA-318', 'LWA-336', 'LWA-335', 'LWA-343', 'LWA-337',
           'LWA-351', 'LWA-344', 'LWA-360', 'LWA-354', 'LWA-002', 'LWA-001',
           'LWA-004', 'LWA-003', 'LWA-006', 'LWA-005', 'LWA-009', 'LWA-007',
           'LWA-011', 'LWA-010', 'LWA-012', 'LWA-008', 'LWA-040', 'LWA-038',
           'LWA-042', 'LWA-041', 'LWA-044', 'LWA-043', 'LWA-046', 'LWA-045',
           'LWA-071', 'LWA-047', 'LWA-074', 'LWA-073', 'LWA-077', 'LWA-075',
           'LWA-275', 'LWA-274', 'LWA-302', 'LWA-286', 'LWA-363', 'LWA-323',
           'LWA-013', 'LWA-016', 'LWA-015', 'LWA-014', 'LWA-018', 'LWA-017',
           'LWA-020', 'LWA-019', 'LWA-022', 'LWA-021', 'LWA-024', 'LWA-023',
           'LWA-026', 'LWA-025', 'LWA-029', 'LWA-027', 'LWA-049', 'LWA-048',
           'LWA-051', 'LWA-050', 'LWA-053', 'LWA-052', 'LWA-054', 'LWA-080',
           'LWA-084', 'LWA-055', 'LWA-324', 'LWA-262', 'LWA-348', 'LWA-331',
           'LWA-365', 'LWA-364', 'LWA-030', 'LWA-028', 'LWA-032', 'LWA-031',
           'LWA-063', 'LWA-060', 'LWA-057', 'LWA-056', 'LWA-059', 'LWA-058',
           'LWA-062', 'LWA-061', 'LWA-085', 'LWA-064', 'LWA-087', 'LWA-086',
           'LWA-089', 'LWA-096', 'LWA-091', 'LWA-090', 'LWA-093', 'LWA-092',
           'LWA-095', 'LWA-094', 'LWA-122', 'LWA-121', 'LWA-325', 'LWA-316',
           'LWA-334', 'LWA-328', 'LWA-361', 'LWA-358', 'LWA-035', 'LWA-033',
           'LWA-034', 'LWA-036', 'LWA-066', 'LWA-065', 'LWA-068', 'LWA-067',
           'LWA-070', 'LWA-069', 'LWA-097', 'LWA-072', 'LWA-099', 'LWA-098',
           'LWA-101', 'LWA-100', 'LWA-103', 'LWA-102', 'LWA-105', 'LWA-104',
           'LWA-129', 'LWA-037', 'LWA-131', 'LWA-130', 'LWA-134', 'LWA-132',
           'LWA-252', 'LWA-139', 'LWA-294', 'LWA-289', 'LWA-319', 'LWA-299',
           'LWA-078', 'LWA-076', 'LWA-081', 'LWA-079', 'LWA-108', 'LWA-107',
           'LWA-110', 'LWA-109', 'LWA-112', 'LWA-111', 'LWA-114', 'LWA-113',
           'LWA-116', 'LWA-115', 'LWA-118', 'LWA-117', 'LWA-120', 'LWA-082',
           'LWA-143', 'LWA-142', 'LWA-145', 'LWA-144', 'LWA-147', 'LWA-150',
           'LWA-149', 'LWA-148', 'LWA-151', 'LWA-172', 'LWA-279', 'LWA-178',
           'LWA-355', 'LWA-349', 'LWA-124', 'LWA-127', 'LWA-126', 'LWA-125',
           'LWA-188', 'LWA-128', 'LWA-153', 'LWA-152', 'LWA-155', 'LWA-154',
           'LWA-157', 'LWA-156', 'LWA-159', 'LWA-158', 'LWA-182', 'LWA-160',
           'LWA-185', 'LWA-184', 'LWA-187', 'LWA-186', 'LWA-190', 'LWA-189',
           'LWA-191', 'LWA-192', 'LWA-224', 'LWA-222', 'LWA-326', 'LWA-322',
           'LWA-345', 'LWA-333', 'LWA-359', 'LWA-347', 'LWA-135', 'LWA-133',
           'LWA-137', 'LWA-136', 'LWA-140', 'LWA-138', 'LWA-161', 'LWA-141',
           'LWA-163', 'LWA-162', 'LWA-165', 'LWA-164', 'LWA-167', 'LWA-166',
           'LWA-193', 'LWA-226', 'LWA-195', 'LWA-194', 'LWA-197', 'LWA-196',
           'LWA-200', 'LWA-199', 'LWA-202', 'LWA-201', 'LWA-227', 'LWA-225',
           'LWA-287', 'LWA-253', 'LWA-293', 'LWA-290', 'LWA-357', 'LWA-329',
           'LWA-170', 'LWA-234', 'LWA-173', 'LWA-171', 'LWA-176', 'LWA-175',
           'LWA-204', 'LWA-203', 'LWA-206', 'LWA-205', 'LWA-208', 'LWA-207',
           'LWA-210', 'LWA-209', 'LWA-228', 'LWA-230', 'LWA-231', 'LWA-229',
           'LWA-233', 'LWA-232', 'LWA-236', 'LWA-235', 'LWA-238', 'LWA-237',
           'LWA-240', 'LWA-239', 'LWA-297', 'LWA-254', 'LWA-346', 'LWA-342',
           'LWA-356', 'LWA-350', 'LWA-177', 'LWA-247', 'LWA-180', 'LWA-179',
           'LWA-183', 'LWA-181', 'LWA-212', 'LWA-211', 'LWA-214', 'LWA-213',
           'LWA-243', 'LWA-215', 'LWA-218', 'LWA-217', 'LWA-221', 'LWA-219',
           'LWA-241', 'LWA-223', 'LWA-244', 'LWA-242', 'LWA-246', 'LWA-245',
           'LWA-249', 'LWA-248', 'LWA-251', 'LWA-250', 'LWA-261', 'LWA-256',
           'LWA-300', 'LWA-264', 'LWA-315', 'LWA-304']
interm = [_smart_int(v.replace('-', '')) for v in interim]
## Sort by swapping until there is nothing left to swap
while True:
    orig_order = [ant.id for ant in ovro.antennas]
    done = True
    for i,j in enumerate(interm):
        k = orig_order.index(j)
        if i != k:
            temp = ovro.antennas[k]
            ovro.antennas[k] = ovro.antennas[i]
            ovro.antennas[i] = temp
            done = False
            break
    if done:
        break
## Trim
ovro.antennas = ovro.antennas[:352]

import numpy,import astropy,from astropy


from functools import wraps
import logging
import sys
import warnings

if sys.version_info[:2] >= (3, 8):
    from functools import cached_property
else:
    from backports.cached_property import cached_property

import numpy as np
import scipy.signal
import pandas as pd

from endaq.batch import quat
from endaq.batch.utils import ide_utils
from endaq.batch.utils.calc import psd, stats, shock, integrate, filters

def as_series(
    unit_type,
    data_name,
    *,
    edit_axis_names=lambda axis_names: axis_names,
    default_axis_names,
):
    """Format method output as a pandas Series."""

    def decorator(method):
        @wraps(method)
        def wrapper(self, *args, **kwargs):
            data = method(self, *args, **kwargs)

            try:
                ch_struct = self._channels[unit_type]
            except KeyError:
                axis_names = default_axis_names
            else:
                axis_names = edit_axis_names(ch_struct.axis_names)

            return pd.Series(
                data,
                index=pd.Index(axis_names, name="axis"),
                name=data_name,
            )

        return wrapper

    return decorator

class Analyzer:
    """
    A class which will run the analyses for the endaq cloud work.  Take the
    important info from the IDE document so that it can be released from memory.
    """

    # TODO: This needs to be reworked, it's currently not very efficient.
    #       -
    #       The current strategy is to:
    #         1. Create a list of time indices for each segment, where the
    #            beginning of the first segment is the earliest data present
    #            among all channels (not al channels start and end at the same time)
    #         2. Create a numpy array from each channel, and grab metadata such
    #            as sampling frequency.
    #         3. For channels with a low sampling frequency, resample them to
    #            five times the segment frequency.
    #         4. Pad each channel with nans so that data exists for each channel
    #            in each time segment.
    #         5. Split each of these by the times created in step 1
    #         6. Lazily run analyses as they're called and cache results.
    #       -
    #       The issue is that this is somewhat memory intensive and just doesn't
    #       need to be. I think we can more efficiently do piecewise evaluation
    #       and re-use intermediate calculations without saving all of the file
    #       as an array.
    #       _
    #       Proposed strategy:
    #         1. Same as above, create a list of time indices for each segment,
    #            where the beginning of the first segment is the earliest data
    #            present among all channels (not al channels start and end at
    #            the same time)
    #         2. For each sensor type:
    #              a. If the sensor is not present, create a list of NaNs in the
    #                 correct shape.
    #              b. For sensors with low sampling rates, use some sort of
    #                 wrapper or generator that will still lazily load data from
    #                 the file but will return interpolated data as if it had a
    #                 higher sample rate (five times the segment frequency)
    #              c. Iterate through the segments.  If the segment is out of
    #                 range, skip it and place Nans in each analysis for that
    #                 channel type.
    #              d. Do the analyses for that segment.  This will allow us to
    #                 reuse intermediate calculations, such as the FFT of the
    #                 acceleration.  This is used to calculate the RMS for the
    #                 acceleration and its integrals.
    #

    MPS2_TO_G = 1 / 9.80665
    MPS_TO_MMPS = 1000
    MPS_TO_UMPS = 10 ** 6
    MPS_TO_KMPH = 3600 / 1e3
    M_TO_MM = 1000

    PV_NATURAL_FREQS = np.logspace(0, 12, base=2, num=12 * 12 + 1, endpoint=True)

    def __init__(
        self,
        doc,
        *,
        preferred_chs=[],
        accel_highpass_cutoff,
        accel_start_time,
        accel_end_time,
        accel_start_margin,
        accel_end_margin,
        psd_freq_bin_width,
        psd_window="hanning",
        pvss_init_freq,
        pvss_bins_per_octave,
        vc_init_freq,
        vc_bins_per_octave,
    ):
        """
        Copies out the numpy arrays for the highest priority channel for each
        sensor type, and any relevant metadata.  Cuts them into chunks.
        """
        if accel_start_time is not None and accel_start_margin is not None:
            raise ValueError(
                "only one of `accel_start_time` and `accel_start_margin` may be set at once"
            )
        if accel_end_time is not None and accel_end_margin is not None:
            raise ValueError(
                "only one of `accel_end_time` and `accel_end_margin` may be set at once"
            )

        self._channels = ide_utils.dict_chs_best(
            (
                (utype, ch_struct)
                for (utype, ch_struct) in ide_utils.chs_by_utype(doc)
                if len(ch_struct.eventarray) > 0
            ),
            max_key=lambda x: (x.channel.id in preferred_chs, len(x.eventarray)),
        )

        self._filename = doc.filename
        self._accelerationFs = None  # gets set in `_accelerationData`
        self._accel_highpass_cutoff = accel_highpass_cutoff
        self._accel_start_time = accel_start_time
        self._accel_end_time = accel_end_time
        self._accel_start_margin = accel_start_margin
        self._accel_end_margin = accel_end_margin
        self._psd_window = psd_window
        self._psd_freq_bin_width = psd_freq_bin_width
        self._pvss_init_freq = pvss_init_freq
        self._pvss_bins_per_octave = pvss_bins_per_octave
        self._vc_init_freq = vc_init_freq
        self._vc_bins_per_octave = vc_bins_per_octave

    # ==========================================================================
    # Data Processing, just to make init cleaner
    # ==========================================================================

    @cached_property
    def _accelerationData(self):
        """Populate the _acceleration* fields, including splitting and extending data."""
        ch_struct = self._channels.get("acc", None)
        if ch_struct is None:
            logging.warning(f"no acceleration channel in {self._filename}")
            return np.empty((3, 0), dtype=np.float)

        aUnits = ch_struct.units[1]
        try:
            conversionFactor = {  # core units = m/s^2
                "g": 1 / self.MPS2_TO_G,
                "m/s\u00b2": 1,
            }[aUnits.lower()]
        except KeyError:
            raise ValueError(f'unknown acceleration channel units "{aUnits}"')

        self._accelerationName = ch_struct.channel.name
        self._accelerationFs = ch_struct.fs

        times = ch_struct.eventarray.arraySlice()[0] * 1e-6  # us -> s
        aData = conversionFactor * ch_struct.eventarray.arrayValues(
            subchannels=ch_struct.sch_ids,
        )

        if self._accel_start_margin is not None:
            margin = int(np.ceil(ch_struct.fs * self._accel_start_margin))
            aData = aData[:, margin:]
        elif self._accel_start_time is not None:
            i = np.searchsorted(times, self._accel_start_time)
            aData = aData[:, i:]
        if self._accel_end_margin is not None:
            margin = int(np.ceil(ch_struct.fs * self._accel_end_margin))
            aData = aData[:, : (-margin or None)]
        elif self._accel_end_time is not None:
            i = np.searchsorted(times, self._accel_end_time)
            aData = aData[:, :i]

        if self._accel_highpass_cutoff:
            aData = filters.highpass(
                aData, fs=ch_struct.fs, cutoff=self._accel_highpass_cutoff, axis=-1
            )

        return aData

    @cached_property
    def _accelerationResultant(self):
        return stats.L2_norm(self._accelerationData, axis=0)

    @cached_property
    def _microphoneData(self):
        """Populate the _microphone* fields, including splitting and extending data."""
        ch_struct = self._channels.get("mic", None)
        if ch_struct is None:
            return np.empty(0, dtype=np.float)

        units = ch_struct.units[1]
        if units.lower() != "a":
            raise ValueError(f'unknown microphone channel units "{units}"')

        self._micName = ch_struct.channel.name
        self._micFs = ch_struct.fs
        data = ch_struct.eventarray.arrayValues(subchannels=ch_struct.sch_ids)

        return data

    @cached_property
    def _velocityData(self):
        aData = self._accelerationData
        if aData.size == 0:
            return np.empty((3, 0), dtype=np.float)

        if not self._accel_highpass_cutoff:
            logging.warning(
                "no highpass filter used before integration; "
                "velocity calculation may be unstable"
            )

        vData = integrate._integrate(aData, dt=1 / self._accelerationFs, axis=1)

        return vData

    @cached_property
    def _displacementData(self):
        vData = self._velocityData
        if vData.size == 0:
            return np.empty((3, 0), dtype=np.float)

        if not self._accel_highpass_cutoff:
            logging.warning(
                "no highpass filter used before integration; "
                "displacement calculation may be unstable"
            )

        dData = integrate._integrate(vData, dt=1 / self._accelerationFs, axis=1)

        return dData

    @cached_property
    def _PVSSData(self):
        aData = self._accelerationData
        if aData.size == 0:
            return np.empty(0, dtype=np.float), self._accelerationData

        log2_f0 = np.log2(self._pvss_init_freq)
        log2_f1 = np.log2(self._accelerationFs)
        num_bins = np.floor(
            self._pvss_bins_per_octave * (log2_f1 - 1 - log2_f0)
        ).astype(int)

        freqs = np.logspace(
            start=log2_f0,
            stop=log2_f0 + num_bins / self._pvss_bins_per_octave,
            num=num_bins + 1,
            base=2,
            endpoint=True,
        )
        freqs = freqs[
            (freqs >= self._accelerationFs / self._accelerationData.shape[-1])
        ]
        pv = shock.pseudo_velocity(
            self._accelerationData,
            freqs,
            dt=1 / self._accelerationFs,
            damp=0.05,
            two_sided=False,
            axis=-1,
        )
        assert pv.ndim == 2
        assert 1 <= pv.shape[0] <= 3
        assert pv.shape[-1] == len(freqs)

        return freqs, pv

    @cached_property
    def _PSDData(self):
        aData = self._accelerationData
        if aData.size == 0:
            return np.empty(0, dtype=np.float), self._accelerationData

        return scipy.signal.welch(
            aData,
            fs=self._accelerationFs,
            nperseg=int(np.ceil(self._accelerationFs / self._psd_freq_bin_width)),
            window=self._psd_window,
            average="median",
            axis=1,
        )

    @cached_property
    def _VCCurveData(self):
        """Calculate Vibration Criteria (VC) Curves for the accelerometer."""
        aData = self._accelerationData
        if aData.size == 0:
            return np.empty(0, dtype=np.float), self._accelerationData

        """
        Theory behind the calculation:

        Let x(t) be a real-valued time-domain signal, and X(2πf) = F{x(t)}(2πf)
        be the Fourier Transform of that signal. By Parseval's Theorem,

            ∫x(t)^2 dt = ∫|X(2πf)|^2 df

        (see https://en.wikipedia.org/wiki/Parseval%27s_theorem#Notation_used_in_physics)

        Rewriting the right side of that equation in the discrete form becomes

            ∫x(t)^2 dt ≈ ∑ |X[k]|^2 • ∆f

        where ∆f = fs/N = (1/∆t) / N = 1/T.
        Limiting the right side to a range of discrete frequencies (k_0, k_1):

            ∫x(t)^2 dt ≈ [∑; k=k_0 -> k≤k_1] |X[k]|^2 • ∆f

        The VC curve calculation is the RMS over the time-domain. If T is the
        duration of the time-domain signal, then:

            √((1/T) ∫x(t)^2 dt)
                ≈ √((1/T) [∑; k=k_0 -> k≤k_1] |X[k]|^2 • ∆f)
                = ∆f • √([∑; k=k_0 -> k≤k_1] |X[k]|^2)

        If the time-series data is acceleration, then the signal needs to first
        be integrated into velocity. This can be done in the frequency domain
        by replacing |X(2πf)|^2 with (1/2πf)^2 |X(2πf)|^2.
        """
        f, a_psd = self._PSDData
        f, v_psd = psd.differentiate(f, a_psd, n=-1)
        f_oct, v_psd_oct = psd.to_octave(
            f,
            v_psd,
            fstart=self._vc_init_freq,
            octave_bins=self._vc_bins_per_octave,
            mode="sum",
        )
        v_vc = np.sqrt(f[1] * v_psd_oct)  # the PSD must already scale by ∆f?

        return f_oct, v_vc

    @cached_property
    def _pressureData(self):
        """Populate the _pressure* fields, including splitting and extending data."""
        ch_struct = self._channels.get("pre", None)
        if ch_struct is None:
            return np.empty(0, dtype=np.float)

        units = ch_struct.units[1]
        try:
            conversionFactor = {  # core units = kPa
                "pa": 1e-3,
                "psi": 6.89476,
                "atm": 101.325,
            }[units.lower()]
        except KeyError:
            raise ValueError(f'unknown pressure channel units "{units}"')

        self._preName = ch_struct.channel.name
        self._preFs = ch_struct.fs
        data = conversionFactor * ch_struct.eventarray.arrayValues(
            subchannels=ch_struct.sch_ids,
        )

        return data

    @cached_property
    def _temperatureData(self):
        """Populate the _temperature* fields, including splitting and extending data."""
        ch_struct = self._channels.get("tmp", None)
        if ch_struct is None:
            return np.empty(0, dtype=np.float)

        units = ch_struct.units[1]
        try:
            conversionFactor, conversionOffset = {  # core units = degrees C
                "\xb0c": (1, 0),
                "\xb0k": (1, 273.15),
                "\xb0f": (5 / 9, -32 * (5 / 9)),
            }[units.lower()]
        except KeyError:
            raise ValueError(f'unknown temperature channel units "{units}"')

        self._tmpName = ch_struct.channel.name
        self._tmpFs = ch_struct.fs
        data = conversionOffset + conversionFactor * ch_struct.eventarray.arrayValues(
            subchannels=ch_struct.sch_ids,
        )

        return data

    @cached_property
    def _gyroscopeData(self):
        """Populate the _gyro* fields, including splitting and extending data."""
        ch_struct = self._channels.get("gyr", None)
        if ch_struct is None:
            return np.empty((3, 0), dtype=np.float)

        self._gyroName = ch_struct.channel.name
        self._gyroFs = ch_struct.fs

        units = ch_struct.units[1]
        if units.lower() == "q":
            quat_array = ch_struct.eventarray.arrayValues(subchannels=ch_struct.sch_ids)
            quat_raw = quat_array[
                [3, 0, 1, 2]
            ]  # reorders to <W, X, Y, Z> & strips out the "Acc" channel

            data = (180 / np.pi) * quat.quat_to_angvel(quat_raw.T, 1 / self._gyroFs).T

            def strip_invalid_prefix(data, prefix_len):
                """Search prefix for invalid data and remove it (if any)."""
                data_mag = stats.L2_norm(data[: 4 * prefix_len], axis=0)
                # the derivative method for `quat_to_angvel` uses the *average*
                # of adjacent differences
                # -> any rotation spikes will result in two adjacent,
                # nearly-equal peaks
                data_agg = 0.5 * (data_mag[:-1] + data_mag[1:])
                argmax_prefix = np.argmax(data_agg[:prefix_len])

                # Prefix data is considered "anomalous" if it is much larger
                # than any surrounding data
                if data_agg[argmax_prefix] > 2 * data_mag[prefix_len:].max():
                    data = data[..., argmax_prefix + 2 :]

                return data

            data = strip_invalid_prefix(
                data, prefix_len=max(4, int(np.ceil(0.25 * self._gyroFs)))
            )
        elif units.lower() in ("dps", "deg/s"):
            data = ch_struct.eventarray.arrayValues(subchannels=ch_struct.sch_ids)
        else:
            raise ValueError(f'unknown gyroscope channel units "{units}"')

        return data

    def _processHumidity(self, channel):
        """Populate the _humidity* fields, including splitting and extending data."""
        pass

    @cached_property
    def _gpsPositionData(self):
        ch_struct = self._channels.get("gps", None)
        if ch_struct is None:
            return np.empty((2, 0), dtype=np.float)

        units = ch_struct.units[1]
        if units.lower() != "degrees":
            raise ValueError(f'unknown GPS position channel units "{units}"')

        data = ch_struct.eventarray.arrayValues(subchannels=ch_struct.sch_ids)

        self._gpsName = ch_struct.channel.name
        self._gpsFs = ch_struct.fs
        # resampling destroys last values -> no resampling

        return data

    @cached_property
    def _gpsSpeedData(self):
        ch_struct = self._channels.get("spd", None)
        if ch_struct is None:
            return np.empty(0)

        units = ch_struct.units[1]
        if units != "m/s":
            raise ValueError(f'unknown GPS ground speed channel units "{units}"')

        data = self.MPS_TO_KMPH * ch_struct.eventarray.arrayValues(
            subchannels=ch_struct.sch_ids
        )

        self._gpsSpeedName = ch_struct.channel.name
        self._gpsSpeedFs = ch_struct.fs

        return data

    # ==========================================================================
    # Analyses
    # ==========================================================================

    @cached_property
    @as_series(
        "acc",
        "RMS Acceleration",
        edit_axis_names=lambda axis_names: axis_names + ["Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def accRMSFull(self):
        """Accelerometer Tri-axial RMS."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            rms = stats.rms(
                self._accelerationData, axis=1
            )  # RuntimeWarning: Mean of empty slice.
        return self.MPS2_TO_G * np.append(rms, stats.L2_norm(rms))

    @cached_property
    @as_series(
        "acc",
        "RMS Velocity",
        edit_axis_names=lambda axis_names: axis_names + ["Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def velRMSFull(self):
        """Velocity Tri-axial RMS, after applying a 0.1Hz highpass filter."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            rms = stats.rms(
                self._velocityData, axis=1
            )  # RuntimeWarning: Mean of empty slice.
        return self.MPS_TO_MMPS * np.append(rms, stats.L2_norm(rms))

    @cached_property
    @as_series(
        "acc",
        "RMS Displacement",
        edit_axis_names=lambda axis_names: axis_names + ["Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def disRMSFull(self):
        """Displacement Tri-axial RMS, after applying a 0.1Hz highpass filter."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            rms = stats.rms(
                self._displacementData, axis=1
            )  # RuntimeWarning: Mean of empty slice.
        return self.M_TO_MM * np.append(rms, stats.L2_norm(rms))

    @cached_property
    @as_series(
        "acc",
        "Peak Absolute Acceleration",
        edit_axis_names=lambda axis_names: axis_names + ["Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def accPeakFull(self):
        """Peak instantaneous tri-axial acceleration"""
        max_abs = stats.max_abs(self._accelerationData, axis=1)
        max_abs_res = np.amax(
            stats.L2_norm(self._accelerationData, axis=0),
            initial=-np.inf,
            axis=-1,
        )
        return self.MPS2_TO_G * np.nan_to_num(
            np.append(max_abs, max_abs_res), nan=np.nan, posinf=np.inf, neginf=np.nan
        )

    @cached_property
    @as_series(
        "acc",
        "Peak Pseudo Velocity Shock Spectrum",
        edit_axis_names=lambda axis_names: axis_names + ["Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def pseudoVelPeakFull(self):
        """Peak Pseudo Velocity"""
        if self._PVSSData[1].size == 0:
            return np.full(self._PVSSData[1].shape[0] + 1, np.nan)

        pv = self._PVSSData[1]
        max_pv = np.amax(pv, initial=-np.inf, axis=1)
        max_pv_res = np.amax(stats.L2_norm(pv, axis=0), initial=-np.inf, axis=-1)
        return self.MPS_TO_MMPS * np.nan_to_num(
            np.append(max_pv, max_pv_res), nan=np.nan, posinf=np.inf, neginf=np.nan
        )

    @cached_property
    @as_series("gps", "GPS Position", default_axis_names=["Latitude", "Longitude"])
    def gpsLocFull(self):
        """Average GPS location"""
        data = self._gpsPositionData
        # 0's occur when gps doesn't have a "lock" -> remove them
        data = data[:, np.all(data != 0, axis=0)]
        if data.size == 0:
            return [np.nan, np.nan]

        return data[..., -1]

    @cached_property
    @as_series("spd", "GPS Speed", default_axis_names=["Ground"])
    def gpsSpeedFull(self):
        """Average GPS speed"""
        data = self._gpsSpeedData
        # 0's occur when gps doesn't have a "lock" -> remove them
        data = data[data != 0]

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return np.mean(data)  # RuntimeWarning: Mean of empty slice.

    @cached_property
    @as_series(
        "gyr",
        "RMS Angular Velocity",
        edit_axis_names=lambda axis_names: ["X", "Y", "Z", "Resultant"],
        default_axis_names=["X", "Y", "Z", "Resultant"],
    )
    def gyroRMSFull(self):
        """Gyroscope RMS"""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            rms = stats.rms(
                self._gyroscopeData, axis=1
            )  # RuntimeWarning: Mean of empty slice.
        return np.append(rms, stats.L2_norm(rms))

    @cached_property
    @as_series("mic", "RMS Microphone", default_axis_names=[""])
    def micRMSFull(self):
        """Microphone RMS"""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return stats.rms(
                self._microphoneData
            )  # RuntimeWarning: Mean of empty slice.

    @cached_property
    @as_series("tmp", "Average Temperature", default_axis_names=[""])
    def tempFull(self):
        """Average Temperature"""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return self._temperatureData.mean()  # RuntimeWarning: Mean of empty slice.

    @cached_property
    @as_series("pre", "Average Pressure", default_axis_names=[""])
    def pressFull(self):
        """Average Pressure"""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return self._pressureData.mean()  # RuntimeWarning: Mean of empty slice.


import io as sysio
import time

import numba
import numpy as np
from scipy.interpolate import interp1d

from det3d.ops.nms.nms_gpu import rotate_iou_gpu_eval
from det3d.core.bbox import box_np_ops
from det3d.datasets.utils.eval import box3d_overlap_kernel
from det3d.datasets.utils.eval import box3d_overlap
from det3d.datasets.utils.eval import calculate_iou_partly
from det3d.datasets.utils.eval import prepare_data
from det3d.datasets.utils.eval import compute_statistics_jit

@numba.jit
def get_thresholds(scores: np.ndarray, num_gt, num_sample_pts=41):
    scores.sort()
    scores = scores[::-1]
    current_recall = 0
    thresholds = []
    for i, score in enumerate(scores):
        l_recall = (i + 1) / num_gt
        if i < (len(scores) - 1):
            r_recall = (i + 2) / num_gt
        else:
            r_recall = l_recall
        if ((r_recall - current_recall) < (current_recall - l_recall)) and (
            i < (len(scores) - 1)
        ):
            continue
        # recall = l_recall
        thresholds.append(score)
        current_recall += 1 / (num_sample_pts - 1.0)
    # print(len(thresholds), len(scores), num_gt)
    return thresholds

def clean_data(gt_anno, dt_anno, current_class, difficulty):
    CLASS_NAMES = [
        "car",
        "pedestrian",
        "bicycle",
        "truck",
        "bus",
        "trailer",
        "construction_vehicle",
        "motorcycle",
        "barrier",
        "traffic_cone",
        "cyclist",
    ]
    MIN_HEIGHT = [40, 25, 25]
    MAX_OCCLUSION = [0, 1, 2]
    MAX_TRUNCATION = [0.15, 0.3, 0.5]
    dc_bboxes, ignored_gt, ignored_dt = [], [], []
    current_cls_name = CLASS_NAMES[current_class].lower()
    num_gt = len(gt_anno["name"])
    num_dt = len(dt_anno["name"])
    num_valid_gt = 0
    for i in range(num_gt):
        bbox = gt_anno["bbox"][i]
        gt_name = gt_anno["name"][i].lower()
        height = bbox[3] - bbox[1]
        valid_class = -1
        if gt_name == current_cls_name:
            valid_class = 1
        elif (
            current_cls_name == "Pedestrian".lower()
            and "Person_sitting".lower() == gt_name
        ):
            valid_class = 0
        elif current_cls_name == "Car".lower() and "Van".lower() == gt_name:
            valid_class = 0
        else:
            valid_class = -1
        ignore = False
        if (
            (gt_anno["occluded"][i] > MAX_OCCLUSION[difficulty])
            or (gt_anno["truncated"][i] > MAX_TRUNCATION[difficulty])
            or (height <= MIN_HEIGHT[difficulty])
        ):
            ignore = True
        if valid_class == 1 and not ignore:
            ignored_gt.append(0)
            num_valid_gt += 1
        elif valid_class == 0 or (ignore and (valid_class == 1)):
            ignored_gt.append(1)
        else:
            ignored_gt.append(-1)
        # for i in range(num_gt):
        if (gt_anno["name"][i] == "DontCare") or (gt_anno["name"][i] == "ignore"):
            dc_bboxes.append(gt_anno["bbox"][i])
    for i in range(num_dt):
        if dt_anno["name"][i].lower() == current_cls_name:
            valid_class = 1
        else:
            valid_class = -1
        height = abs(dt_anno["bbox"][i, 3] - dt_anno["bbox"][i, 1])
        if height < MIN_HEIGHT[difficulty]:
            ignored_dt.append(1)
        elif valid_class == 1:
            ignored_dt.append(0)
        else:
            ignored_dt.append(-1)

    return num_valid_gt, ignored_gt, ignored_dt, dc_bboxes

def get_split_parts(num, num_part):
    same_part = num // num_part
    remain_num = num % num_part
    if remain_num == 0:
        return [same_part] * num_part
    else:
        return [same_part] * num_part + [remain_num]

@numba.jit(nopython=True)
def fused_compute_statistics(
    overlaps,
    pr,
    gt_nums,
    dt_nums,
    dc_nums,
    gt_datas,
    dt_datas,
    dontcares,
    ignored_gts,
    ignored_dets,
    metric,
    min_overlap,
    thresholds,
    compute_aos=False,
):
    gt_num = 0
    dt_num = 0
    dc_num = 0
    for i in range(gt_nums.shape[0]):
        for t, thresh in enumerate(thresholds):
            overlap = overlaps[
                dt_num : dt_num + dt_nums[i], gt_num : gt_num + gt_nums[i]
            ]

            gt_data = gt_datas[gt_num : gt_num + gt_nums[i]]
            dt_data = dt_datas[dt_num : dt_num + dt_nums[i]]
            ignored_gt = ignored_gts[gt_num : gt_num + gt_nums[i]]
            ignored_det = ignored_dets[dt_num : dt_num + dt_nums[i]]
            dontcare = dontcares[dc_num : dc_num + dc_nums[i]]
            tp, fp, fn, similarity, _ = compute_statistics_jit(
                overlap,
                gt_data,
                dt_data,
                ignored_gt,
                ignored_det,
                dontcare,
                metric,
                min_overlap=min_overlap,
                thresh=thresh,
                compute_fp=True,
                compute_aos=compute_aos,
            )
            pr[t, 0] += tp
            pr[t, 1] += fp
            pr[t, 2] += fn
            if similarity != -1:
                pr[t, 3] += similarity
        gt_num += gt_nums[i]
        dt_num += dt_nums[i]
        dc_num += dc_nums[i]

def eval_class_v3(
    gt_annos,
    dt_annos,
    current_classes,
    difficultys,
    metric,
    min_overlaps,
    compute_aos=False,
    z_axis=1,
    z_center=1.0,
    num_parts=50,
):
    """Kitti eval. support 2d/bev/3d/aos eval. support 0.5:0.05:0.95 coco AP.
    Args:
        gt_annos: dict, must from get_label_annos() in kitti_common.py
        dt_annos: dict, must from get_label_annos() in kitti_common.py
        current_class: int, 0: car, 1: pedestrian, 2: cyclist
        difficulty: int. eval difficulty, 0: easy, 1: normal, 2: hard
        metric: eval type. 0: bbox, 1: bev, 2: 3d
        min_overlap: float, min overlap. official:
            [[0.7, 0.5, 0.5], [0.7, 0.5, 0.5], [0.7, 0.5, 0.5]]
            format: [metric, class]. choose one from matrix above.
        num_parts: int. a parameter for fast calculate algorithm

    Returns:
        dict of recall, precision and aos
    """
    assert len(gt_annos) == len(dt_annos)
    num_examples = len(gt_annos)
    split_parts = get_split_parts(num_examples, num_parts)
    split_parts = [i for i in split_parts if i != 0]

    rets = calculate_iou_partly(
        dt_annos, gt_annos, metric, num_parts, z_axis=z_axis, z_center=z_center
    )
    overlaps, parted_overlaps, total_dt_num, total_gt_num = rets
    N_SAMPLE_PTS = 41
    num_minoverlap = len(min_overlaps)
    num_class = len(current_classes)
    num_difficulty = len(difficultys)
    precision = np.zeros([num_class, num_difficulty, num_minoverlap, N_SAMPLE_PTS])
    recall = np.zeros([num_class, num_difficulty, num_minoverlap, N_SAMPLE_PTS])
    aos = np.zeros([num_class, num_difficulty, num_minoverlap, N_SAMPLE_PTS])
    all_thresholds = np.zeros([num_class, num_difficulty, num_minoverlap, N_SAMPLE_PTS])
    for m, current_class in enumerate(current_classes):
        for l, difficulty in enumerate(difficultys):
            rets = prepare_data(
                gt_annos,
                dt_annos,
                current_class,
                difficulty=difficulty,
                clean_data=clean_data,
            )
            (
                gt_datas_list,
                dt_datas_list,
                ignored_gts,
                ignored_dets,
                dontcares,
                total_dc_num,
                total_num_valid_gt,
            ) = rets
            for k, min_overlap in enumerate(min_overlaps[:, metric, m]):
                thresholdss = []
                for i in range(len(gt_annos)):
                    rets = compute_statistics_jit(
                        overlaps[i],
                        gt_datas_list[i],
                        dt_datas_list[i],
                        ignored_gts[i],
                        ignored_dets[i],
                        dontcares[i],
                        metric,
                        min_overlap=min_overlap,
                        thresh=0.0,
                        compute_fp=False,
                    )
                    tp, fp, fn, similarity, thresholds = rets
                    thresholdss += thresholds.tolist()
                thresholdss = np.array(thresholdss)
                thresholds = get_thresholds(thresholdss, total_num_valid_gt)
                thresholds = np.array(thresholds)
                # print(thresholds)
                all_thresholds[m, l, k, : len(thresholds)] = thresholds
                pr = np.zeros([len(thresholds), 4])
                idx = 0
                for j, num_part in enumerate(split_parts):
                    gt_datas_part = np.concatenate(
                        gt_datas_list[idx : idx + num_part], 0
                    )
                    dt_datas_part = np.concatenate(
                        dt_datas_list[idx : idx + num_part], 0
                    )
                    dc_datas_part = np.concatenate(dontcares[idx : idx + num_part], 0)
                    ignored_dets_part = np.concatenate(
                        ignored_dets[idx : idx + num_part], 0
                    )
                    ignored_gts_part = np.concatenate(
                        ignored_gts[idx : idx + num_part], 0
                    )
                    fused_compute_statistics(
                        parted_overlaps[j],
                        pr,
                        total_gt_num[idx : idx + num_part],
                        total_dt_num[idx : idx + num_part],
                        total_dc_num[idx : idx + num_part],
                        gt_datas_part,
                        dt_datas_part,
                        dc_datas_part,
                        ignored_gts_part,
                        ignored_dets_part,
                        metric,
                        min_overlap=min_overlap,
                        thresholds=thresholds,
                        compute_aos=compute_aos,
                    )
                    idx += num_part
                for i in range(len(thresholds)):
                    # recall[m, l, k, i] = pr[i, 0] / (pr[i, 0] + pr[i, 2])
                    precision[m, l, k, i] = pr[i, 0] / (pr[i, 0] + pr[i, 1])
                    if compute_aos:
                        aos[m, l, k, i] = pr[i, 3] / (pr[i, 0] + pr[i, 1])
                for i in range(len(thresholds)):
                    precision[m, l, k, i] = np.max(precision[m, l, k, i:], axis=-1)
                    if compute_aos:
                        aos[m, l, k, i] = np.max(aos[m, l, k, i:], axis=-1)
                # use interp to calculate recall
                """
                current_recalls = np.linspace(0, 1, 41)
                prec_unique, inds = np.unique(precision[m, l, k], return_index=True)
                current_recalls = current_recalls[inds]
                f = interp1d(prec_unique, current_recalls)
                precs_for_recall = np.linspace(0, 1, 41)
                max_prec = np.max(precision[m, l, k])
                valid_prec = precs_for_recall < max_prec
                num_valid_prec = valid_prec.sum()
                recall[m, l, k, :num_valid_prec] = f(precs_for_recall[valid_prec])
                """
    ret_dict = {
        "recall": recall,  # [num_class, num_difficulty, num_minoverlap, N_SAMPLE_PTS]
        "precision": precision,
        "orientation": aos,
        "thresholds": all_thresholds,
        "min_overlaps": min_overlaps,
    }
    return ret_dict

def get_mAP2(prec):
    sums = 0
    interval = 4
    for i in range(0, prec.shape[-1], interval):
        sums = sums + prec[..., i]
    return sums / int(prec.shape[-1] / interval) * 100

def get_mAP(prec):
    sums = 0
    for i in range(0, prec.shape[-1], 4):
        sums = sums + prec[..., i]
    return sums / 11 * 100

#def get_mAP(prec):
#    sums = 0
#    for i in range(0, prec.shape[-1], 1):
#        sums = sums + prec[..., i]
#    return sums / 40 * 100

def do_eval_v2(
    gt_annos,
    dt_annos,
    current_classes,
    min_overlaps,
    compute_aos=False,
    difficultys=(0, 1, 2),
    z_axis=1,
    z_center=1.0,
):
    # min_overlaps: [num_minoverlap, metric, num_class]
    ret = eval_class_v3(
        gt_annos,
        dt_annos,
        current_classes,
        difficultys,
        0,
        min_overlaps,
        compute_aos,
        z_axis=z_axis,
        z_center=z_center,
    )
    # ret: [num_class, num_diff, num_minoverlap, num_sample_points]
    mAP_bbox = get_mAP(ret["precision"])
    mAP_aos = None
    if compute_aos:
        mAP_aos = get_mAP(ret["orientation"])
    ret = eval_class_v3(
        gt_annos,
        dt_annos,
        current_classes,
        difficultys,
        1,
        min_overlaps,
        z_axis=z_axis,
        z_center=z_center,
    )
    mAP_bev = get_mAP(ret["precision"])
    ret = eval_class_v3(
        gt_annos,
        dt_annos,
        current_classes,
        difficultys,
        2,
        min_overlaps,
        z_axis=z_axis,
        z_center=z_center,
    )
    mAP_3d = get_mAP(ret["precision"])
    return mAP_bbox, mAP_bev, mAP_3d, mAP_aos

def do_eval_v3(
    gt_annos,
    dt_annos,
    current_classes,
    min_overlaps,
    compute_aos=False,
    difficultys=(0, 1, 2),
    z_axis=1,
    z_center=1.0,
):
    # min_overlaps: [num_minoverlap, metric, num_class]
    types = ["bbox", "bev", "3d"]
    metrics = {}
    for i in range(3):
        ret = eval_class_v3(
            gt_annos,
            dt_annos,
            current_classes,
            difficultys,
            i,
            min_overlaps,
            compute_aos,
            z_axis=z_axis,
            z_center=z_center,
        )
        metrics[types[i]] = ret
    return metrics

def do_coco_style_eval(
    gt_annos,
    dt_annos,
    current_classes,
    overlap_ranges,
    compute_aos,
    z_axis=1,
    z_center=1.0,
):
    # overlap_ranges: [range, metric, num_class]
    min_overlaps = np.zeros([10, *overlap_ranges.shape[1:]])
    for i in range(overlap_ranges.shape[1]):
        for j in range(overlap_ranges.shape[2]):
            min_overlaps[:, i, j] = np.linspace(*overlap_ranges[:, i, j])
    mAP_bbox, mAP_bev, mAP_3d, mAP_aos = do_eval_v2(
        gt_annos,
        dt_annos,
        current_classes,
        min_overlaps,
        compute_aos,
        z_axis=z_axis,
        z_center=z_center,
    )
    # ret: [num_class, num_diff, num_minoverlap]
    mAP_bbox = mAP_bbox.mean(-1)
    mAP_bev = mAP_bev.mean(-1)
    mAP_3d = mAP_3d.mean(-1)
    if mAP_aos is not None:
        mAP_aos = mAP_aos.mean(-1)
    return mAP_bbox, mAP_bev, mAP_3d, mAP_aos

def print_str(value, *arg, sstream=None):
    if sstream is None:
        sstream = sysio.StringIO()
    sstream.truncate(0)
    sstream.seek(0)
    print(value, *arg, file=sstream)
    return sstream.getvalue()

def get_official_eval_result(
    gt_annos, dt_annos, current_classes, difficultys=[0, 1, 2], z_axis=1, z_center=1.0
):
    """
        gt_annos and dt_annos must contains following keys:
        [bbox, location, dimensions, rotation, score]
    """
    overlap_mod = np.array(
        [
            [0.7, 0.5, 0.5, 0.7, 0.7, 0.7, 0.7, 0.5, 0.5, 0.5, 0.5],
            [0.7, 0.5, 0.5, 0.7, 0.7, 0.7, 0.7, 0.5, 0.5, 0.5, 0.5],
            [0.7, 0.5, 0.5, 0.7, 0.7, 0.7, 0.7, 0.5, 0.5, 0.5, 0.5],
        ]
    )
    overlap_easy = np.array(
        [
            [0.7, 0.5, 0.5, 0.7, 0.7, 0.7, 0.7, 0.5, 0.25, 0.25, 0.5],
            [0.5, 0.25, 0.25, 0.5, 0.5, 0.5, 0.5, 0.25, 0.25, 0.25, 0.25],
            [0.5, 0.25, 0.25, 0.5, 0.5, 0.5, 0.5, 0.25, 0.25, 0.25, 0.25],
        ]
    )
    min_overlaps = np.stack([overlap_mod, overlap_easy], axis=0)  # [2, 3, 5]
    class_to_name = {
        0: "car",
        1: "pedestrian",
        2: "bicycle",
        3: "truck",
        4: "bus",
        5: "trailer",
        6: "construction_vehicle",
        7: "motorcycle",
        8: "barrier",
        9: "traffic_cone",
        10: "cyclist",
    }
    name_to_class = {v: n for n, v in class_to_name.items()}
    if not isinstance(current_classes, (list, tuple)):
        current_classes = [current_classes]
    current_classes_int = []
    for curcls in current_classes:
        if isinstance(curcls, str):
            current_classes_int.append(name_to_class[curcls.lower()])
        else:
            current_classes_int.append(curcls)
    current_classes = current_classes_int
    min_overlaps = min_overlaps[:, :, current_classes]
    result = ""
    # check whether alpha is valid
    compute_aos = False
    for anno in dt_annos:
        if anno["alpha"].shape[0] != 0:
            if anno["alpha"][0] != -10:
                compute_aos = True
            break
    # TODO dt2gt
    metrics = do_eval_v3(
        gt_annos,
        dt_annos,
        current_classes,
        min_overlaps,
        compute_aos,
        difficultys,
        z_axis=z_axis,
        z_center=z_center,
    )
    detail = {}
    for j, curcls in enumerate(current_classes):
        # mAP threshold array: [num_minoverlap, metric, class]
        # mAP result: [num_class, num_diff, num_minoverlap]
        class_name = class_to_name[curcls]
        detail[class_name] = {}
        for i in range(min_overlaps.shape[0]):
            mAPbbox = get_mAP(metrics["bbox"]["precision"][j, :, i])
            mAPbev = get_mAP(metrics["bev"]["precision"][j, :, i])
            mAP3d = get_mAP(metrics["3d"]["precision"][j, :, i])
            detail[class_name][f"bbox@{min_overlaps[i, 0, j]:.2f}"] = mAPbbox.tolist()
            detail[class_name][f"bev@{min_overlaps[i, 1, j]:.2f}"] = mAPbev.tolist()
            detail[class_name][f"3d@{min_overlaps[i, 2, j]:.2f}"] = mAP3d.tolist()

            result += print_str(
                (
                    f"{class_to_name[curcls]} "
                    "AP(Average Precision)@{:.2f}, {:.2f}, {:.2f}:".format(
                        *min_overlaps[i, :, j]
                    )
                )
            )
            mAPbbox = ", ".join(f"{v:.2f}" for v in mAPbbox)
            mAPbev = ", ".join(f"{v:.2f}" for v in mAPbev)
            mAP3d = ", ".join(f"{v:.2f}" for v in mAP3d)
            result += print_str(f"bbox AP:{mAPbbox}")
            result += print_str(f"bev  AP:{mAPbev}")
            result += print_str(f"3d   AP:{mAP3d}")
            if compute_aos:
                mAPaos = get_mAP(metrics["bbox"]["orientation"][j, :, i])
                detail[class_name][f"aos"] = mAPaos.tolist()
                mAPaos = ", ".join(f"{v:.2f}" for v in mAPaos)
                result += print_str(f"aos  AP:{mAPaos}")
    return {
        "result": result,
        "detail": detail,
    }

def get_coco_eval_result(gt_annos, dt_annos, current_classes, z_axis=1, z_center=1.0):
    class_to_name = {
        0: "car",
        1: "pedestrian",
        2: "bicycle",
        3: "truck",
        4: "bus",
        5: "trailer",
        6: "construction_vehicle",
        7: "motorcycle",
        8: "barrier",
        9: "traffic_cone",
        10: "cyclist",
    }
    class_to_range = {
        0: [0.5, 0.95, 10],
        1: [0.25, 0.7, 10],
        2: [0.25, 0.7, 10],
        3: [0.5, 0.95, 10],
        4: [0.5, 0.95, 10],
        5: [0.5, 0.95, 10],
        6: [0.5, 0.95, 10],
        7: [0.25, 0.7, 10],
        8: [0.25, 0.7, 10],
        9: [0.25, 0.7, 10],
        10: [0.25, 0.7, 10],
    }
    # class_to_range = {
    #     0: [0.5, 0.95, 10],
    #     1: [0.25, 0.7, 10],
    #     2: [0.25, 0.7, 10],
    #     3: [0.5, 0.95, 10],
    #     4: [0.25, 0.7, 10],
    #     5: [0.5, 0.95, 10],
    #     6: [0.5, 0.95, 10],
    #     7: [0.5, 0.95, 10],
    # }

    name_to_class = {v: n for n, v in class_to_name.items()}
    if not isinstance(current_classes, (list, tuple)):
        current_classes = [current_classes]
    current_classes_int = []
    for curcls in current_classes:
        if isinstance(curcls, str):
            current_classes_int.append(name_to_class[curcls.lower()])
        else:
            current_classes_int.append(curcls)
    current_classes = current_classes_int
    overlap_ranges = np.zeros([3, 3, len(current_classes)])
    for i, curcls in enumerate(current_classes):
        overlap_ranges[:, :, i] = np.array(class_to_range[curcls])[:, np.newaxis]
    result = ""
    # check whether alpha is valid
    compute_aos = False
    for anno in dt_annos:
        if anno["alpha"].shape[0] != 0:
            if anno["alpha"][0] != -10:
                compute_aos = True
            break
    mAPbbox, mAPbev, mAP3d, mAPaos = do_coco_style_eval(
        gt_annos,
        dt_annos,
        current_classes,
        overlap_ranges,
        compute_aos,
        z_axis=z_axis,
        z_center=z_center,
    )
    detail = {}
    for j, curcls in enumerate(current_classes):
        class_name = class_to_name[curcls]
        detail[class_name] = {}
        # mAP threshold array: [num_minoverlap, metric, class]
        # mAP result: [num_class, num_diff, num_minoverlap]
        o_range = np.array(class_to_range[curcls])[[0, 2, 1]]
        o_range[1] = (o_range[2] - o_range[0]) / (o_range[1] - 1)
        result += print_str(
            (
                f"{class_to_name[curcls]} "
                "coco AP@{:.2f}:{:.2f}:{:.2f}:".format(*o_range)
            )
        )
        result += print_str(
            (
                f"bbox AP:{mAPbbox[j, 0]:.2f}, "
                f"{mAPbbox[j, 1]:.2f}, "
                f"{mAPbbox[j, 2]:.2f}"
            )
        )
        result += print_str(
            (
                f"bev  AP:{mAPbev[j, 0]:.2f}, "
                f"{mAPbev[j, 1]:.2f}, "
                f"{mAPbev[j, 2]:.2f}"
            )
        )
        result += print_str(
            (f"3d   AP:{mAP3d[j, 0]:.2f}, " f"{mAP3d[j, 1]:.2f}, " f"{mAP3d[j, 2]:.2f}")
        )
        detail[class_name][f"bbox"] = mAPbbox[j].tolist()
        detail[class_name][f"bev"] = mAPbev[j].tolist()
        detail[class_name][f"3d"] = mAP3d[j].tolist()

        if compute_aos:
            detail[class_name][f"aos"] = mAPaos[j].tolist()
            result += print_str(
                (
                    f"aos  AP:{mAPaos[j, 0]:.2f}, "
                    f"{mAPaos[j, 1]:.2f}, "
                    f"{mAPaos[j, 2]:.2f}"
                )
            )
    return {
        "result": result,
        "detail": detail,
    }


"""
Computation of the demagnetising field using the Fredkin-Koehler
technique and the infamous magpar method.

Rationale: The previous implementation in FemBemFKSolver (child class
of FemBemDeMagSolver) was kind of a mess. This does the same thing in the same
time with less code. Should be more conducive to further optimisation or as
a template for other techniques like the GCR.

"""
import numpy as np
import dolfin as df
import logging
from aeon import timer, Timer
from finmag.util.consts import mu0
from finmag.native.llg import compute_bem_fk
from finmag.util.meshes import nodal_volume
from finmag.util import helpers, configuration
from finmag.field import Field
from fk_demag_pbc import BMatrixPBC

logger = logging.getLogger('finmag')
fk_timer = Timer()

class FKDemag(object):

    """
    Computation of the demagnetising field using the Fredkin-Koehler hybrid
    FEM/BEM technique.

    Fredkin, D.R. and Koehler, T.R., "`Hybrid method for computing
    demagnetizing fields`_", IEEE Transactions on Magnetics, vol.26, no.2,
    pp.415-417, Mar 1990.

    .. _Hybrid method for computing demagnetizing fields:
       http://ieeexplore.ieee.org/xpls/abs_all.jsp?arnumber=106342

    """

    def __init__(self, name='Demag', thin_film=False, macrogeometry=None,
                 solver_type=None, parameters=None):
        """
        Create a new FKDemag instance.

        The attribute `parameters` is a dict that contains the settings for the
        solvers for the Neumann (potential phi_1) and Laplace (potential phi_2)
        problems.

        Setting the method used by the solvers:
        Change the entries `phi_1_solver` and `phi_2_solver` to a value from
        `df.list_krylov_solver_methods()`. Default is dolfin's default.

        Setting the preconditioners:
        Change the entries `phi_1_preconditioner` and `phi_2_preconditioner` to
        a value from `df.list_krylov_solver_preconditioners()`. Default is
        dolfin's default. There is a set of parameters optimised for thin films
        (cg/ilu followed by default without preconditioner) that can be used by
        passing in the argument 'thin_film` set to True.

        Setting the tolerances:
        Change the existing entries inside `phi_1` and `phi_2` which are
        themselves dicts. You can add new entries to these dicts as well.
        Everything which is understood by `df.KrylovSolver` is valid.

        Allowed values for `solver_type` are 'Krylov','LU' and `None` (the
        latter uses the value set in the .finmagrc file, defaulting to 'Krylov'
        as no value is provided there).

        """
        self.name = name
        self.in_jacobian = False
        default_parameters = {
            'absolute_tolerance': 1e-6,
            'relative_tolerance': 1e-6,
            'maximum_iterations': int(1e4)
        }
        self.parameters = {
            'phi_1_solver': 'default',
            'phi_1_preconditioner': 'default',
            'phi_1': default_parameters,
            'phi_2_solver': 'default',
            'phi_2_preconditioner': 'default',
            'phi_2': default_parameters.copy()
        }
        if parameters is not None:
            for (k, v) in parameters.items():
                logger.debug(
                    "Setting demag solver parameter {}='{}'".format(k, v))
                if k in ['phi_1', 'phi_2']:
                    # Since self.parameters['phi_1'] is a dictionary itself,
                    # only update the keys that are given (and similarly for
                    # 'phi_2').
                    for (k2, v2) in v.items():
                        self.parameters[k][k2] = v2
                else:
                    self.parameters[k] = v
            logger.debug("Demag parameters now: {}".format(self.parameters))

        self.solver_type = solver_type

        if thin_film:
            self.parameters["phi_1_solver"] = "cg"
            self.parameters["phi_1_preconditioner"] = "ilu"
            self.parameters["phi_2_preconditioner"] = "none"

        self.macrogeometry = macrogeometry

    @timer.method
    def setup(self, m, Ms, unit_length=1):
        """
        Setup the FKDemag instance. Usually called automatically by the
        Simulation object.

        *Arguments*

        m: finmag.Field

            The unit magnetisation on a finite element space.

        Ms: float

            The saturation magnetisation in A/m.

        unit_length: float

            The length (in m) represented by one unit on the mesh. Default 1.

        """
        assert isinstance(m, Field)
        assert isinstance(Ms, Field)

        self.m = m
        self.Ms = Ms
        self.unit_length = unit_length
        self.S1 = df.FunctionSpace(self.m.mesh(), "Lagrange", 1)

        self._test1 = df.TestFunction(self.S1)
        self._trial1 = df.TrialFunction(self.S1)
        self._test3 = df.TestFunction(self.m.functionspace)
        self._trial3 = df.TrialFunction(self.m.functionspace)

        # for computation of energy
        self._nodal_volumes = nodal_volume(self.S1, unit_length)
        self._H_func = df.Function(m.functionspace)  # we will copy field into
        # this when we need the
        # energy
        self._E_integrand = -0.5 * mu0 * \
            df.dot(self._H_func, self.m.f * self.Ms.f)
        self._E = self._E_integrand * df.dx
        self._nodal_E = df.dot(self._E_integrand, self._test1) * df.dx
        self._nodal_E_func = df.Function(self.S1)

        # for computation of field and scalar magnetic potential
        self._poisson_matrix = self._poisson_matrix()
        self._laplace_zeros = df.Function(self.S1).vector()

        # determine the solver type to be used (Krylov or LU); if the kwarg
        # 'solver_type' is not provided, try to read the setting from the
        # .finmagrc file; use 'Krylov' if this fails.
        solver_type = self.solver_type
        if solver_type is None:
            solver_type = configuration.get_config_option(
                'demag', 'solver_type', 'Krylov')
        if solver_type == 'None':  # if the user set 'solver_type = None' in
                                   # the .finmagrc file, solver_type will be a
                                   # string so we need to catch this here.
            solver_type = 'Krylov'
        logger.debug("Using {} solver for demag.".format(solver_type))

        if solver_type == 'Krylov':
            self._poisson_solver = df.KrylovSolver(self._poisson_matrix.copy(),
                                                   self.parameters['phi_1_solver'], self.parameters['phi_1_preconditioner'])
            self._poisson_solver.parameters.update(self.parameters['phi_1'])
            self._laplace_solver = df.KrylovSolver(
                self.parameters['phi_2_solver'], self.parameters['phi_2_preconditioner'])
            self._laplace_solver.parameters.update(self.parameters['phi_2'])
            # We're setting 'same_nonzero_pattern=True' to enforce the
            # same matrix sparsity pattern across different demag solves,
            # which should speed up things.
            #self._laplace_solver.parameters["preconditioner"][
            #    "structure"] = "same_nonzero_pattern"
        elif solver_type == 'LU':
            self._poisson_solver = df.LUSolver(self._poisson_matrix.copy())
            self._laplace_solver = df.LUSolver()
            self._poisson_solver.parameters["reuse_factorization"] = True
            self._laplace_solver.parameters["reuse_factorization"] = True
        else:
            raise ValueError("Argument 'solver_type' must be either 'Krylov' or 'LU'. "
                             "Got: '{}'".format(solver_type))

        with fk_timer('compute BEM'):
            if not hasattr(self, "_bem"):
                if self.macrogeometry is not None:
                    Ts = self.macrogeometry.compute_Ts(self.m.mesh())
                    pbc = BMatrixPBC(self.m.mesh(), Ts)
                    self._b2g_map = np.array(pbc.b2g_map, dtype=np.int)
                    self._bem = pbc.bm
                else:
                    self._bem, self._b2g_map = compute_bem_fk(
                        df.BoundaryMesh(self.m.mesh(), 'exterior', False))
        logger.debug("Boundary element matrix uses {:.2f} MB of memory.".format(
            self._bem.nbytes / 1024. ** 2))
        # solution of inhomogeneous Neumann problem
        self._phi_1 = df.Function(self.S1)
        # solution of Laplace equation inside domain
        self._phi_2 = df.Function(self.S1)
        self._phi = df.Function(self.S1)  # magnetic potential phi_1 + phi_2

        # To be applied to the vector field m as first step of computation of
        # _phi_1.  This gives us div(M), which is equal to Laplace(_phi_1),
        # equation which is then solved using _poisson_solver.
        self._Ms_times_divergence = df.assemble(
            self.Ms.f * df.inner(self._trial3, df.grad(self._test1)) * df.dx)

        # we move the boundary condition here to avoid create a instance each
        # time when compute the magnetic potential
        self.boundary_condition = df.DirichletBC(
            self.S1, self._phi_2, df.DomainBoundary())
        self.boundary_condition.apply(self._poisson_matrix)

        self._setup_gradient_computation()

    @timer.method
    def precomputed_bem(self, bem, b2g_map):
        """
        If the BEM and a boundary to global vertices map are known, they can be
        passed to the FKDemag object with this method so it will skip
        re-computing them.

        """
        self._bem, self._b2g_map = bem, b2g_map

    @timer.method
    def compute_potential(self):
        """
        Compute the magnetic potential.

        *Returns*
            df.Function
                The magnetic potential.

        """
        self._compute_magnetic_potential()
        return self._phi

    @timer.method
    def compute_field(self):
        """
        Compute the demagnetising field.

        *Returns*
            numpy.ndarray
                The demagnetising field.

        """
        self._compute_magnetic_potential()
        return self._compute_gradient()

    def average_field(self):
        """
        Compute the average demag field.
        """
        return helpers.average_field(self.compute_field())

    @timer.method
    def compute_energy(self):
        """
        Compute the total energy of the field.

        .. math::

            E_\\mathrm{d} = -\\frac12 \\mu_0 \\int_\\Omega
            H_\\mathrm{d} \\cdot \\vec M \\mathrm{d}x

        *Returns*
            Float
                The energy of the demagnetising field.

        """
        self._H_func.vector()[:] = self.compute_field()
        return df.assemble(self._E) * self.unit_length ** self.m.mesh_dim()

    @timer.method
    def energy_density(self):
        """
        Compute the energy density in the field.

        .. math::
            \\rho = \\frac{E_{\\mathrm{d}, i}}{V_i},

        where V_i is the volume associated with the node i.

        *Returns*
            numpy.ndarray
                The energy density of the demagnetising field.

        """
        self._H_func.vector()[:] = self.compute_field()
        nodal_E = df.assemble(self._nodal_E).array() * \
            self.unit_length ** self.m.mesh_dim()
        return nodal_E / self._nodal_volumes

    @timer.method
    def energy_density_function(self):
        """
        Returns the energy density in the field as a dolfin function to allow probing.

        *Returns*
            dolfin.Function
                The energy density of the demagnetising field.

        """
        self._nodal_E_func.vector()[:] = self.energy_density()
        return self._nodal_E_func

    @fk_timer.method
    def _poisson_matrix(self):
        A = df.dot(df.grad(self._trial1), df.grad(self._test1)) * df.dx
        return df.assemble(A)  # stiffness matrix for Poisson equation

    def _compute_magnetic_potential(self):
        # compute _phi_1 on the whole domain
        g_1 = self._Ms_times_divergence * self.m.f.vector()
        with fk_timer("first linear solve"):
            self._poisson_solver.solve(self._phi_1.vector(), g_1)

        # compute _phi_2 on the boundary using the Dirichlet boundary
        # conditions we get from BEM * _phi_1 on the boundary.
        with fk_timer("using boundary conditions"):
            phi_1 = self._phi_1.vector()[self._b2g_map]
            self._phi_2.vector()[self._b2g_map[:]] = np.dot(
                self._bem, phi_1)
            #boundary_condition = df.DirichletBC(self.S1, self._phi_2, df.DomainBoundary())
            #A = self._poisson_matrix.copy()
            #b = self._laplace_zeros
            #boundary_condition.apply(A, b)
            A = self._poisson_matrix
            b = self._laplace_zeros
            self.boundary_condition.set_value(self._phi_2)
            self.boundary_condition.apply(A, b)

        # compute _phi_2 on the whole domain
        with fk_timer("second linear solve"):
            self._laplace_solver.solve(A, self._phi_2.vector(), b)

        # add _phi_1 and _phi_2 to obtain magnetic potential
        self._phi.vector()[:] = self._phi_1.vector() + self._phi_2.vector()

    @fk_timer.method
    def _setup_gradient_computation(self):
        """
        Prepare the discretised gradient to use in :py:meth:`FKDemag._compute_gradient`.

        We don't need the gradient field as a continuous field, we are only
        interested in the values at specific points. It is thus a waste of
        computational effort to use a projection of the gradient field, since
        it performs the fairly large operation of assembling a matrix and
        solving a linear system of equations.

        """
        A = df.inner(self._test3, - df.grad(self._trial1)) * df.dx
        # This can be applied to scalar functions.
        self._gradient = df.assemble(A)

        # The `A` above is in fact not quite the gradient, since we integrated
        # over the volume as well. We will divide by the volume later, after
        # the multiplication of the scalar magnetic potential. Since the two
        # operations are symmetric (multiplying by volume, dividing by volume)
        # we don't have to care for the units, i.e. unit_length.
        b = df.dot(self._test3, df.Constant((1, 1, 1))) * df.dx
        self._nodal_volumes_S3_no_units = df.assemble(b).array()

    @fk_timer.method
    def _compute_gradient(self):
        """
        Get the demagnetising field from the magnetic scalar potential.

        .. math::

            \\vec{H}_{\\mathrm{d}} = - \\nabla \\phi (\\vec{r})

        Using dolfin, we would translate this to

        .. sourcecode::

            H_d = df.project(- df.grad(self._phi), self.m.functionspace)

        but the method used here is computationally less expensive.

        """
        H = self._gradient * self._phi.vector()
        return H.array() / self._nodal_volumes_S3_no_units