#! /usr/bin/python3
# -*- coding: utf-8 -*-

"""
    TO DO:
    map to healpix grid
"""

__author__ = 'Alan Loh, Julien Girard'
__copyright__ = 'Copyright 2019, nenupytv'
__credits__ = ['Alan Loh', 'Julien Girard']
__maintainer__ = 'Alan'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'Production'
__all__ = [
    'Grid',
    'Grid_Simple',
    'Grid_HPX'
    ]

# import healpy as hp
import numpy as np
import astropy.units as un

from nenupytv.image import AAFilter

# ============================================================= #
# ---------------------------- Grid --------------------------- #
# ============================================================= #
class Grid(object):
    """ Gridding class

        :param vis: Visibilities
        :type vis: `~np.ndarray`
        :param uvw: UVW coordinates
        :type uvw: `~np.ndarray`
        :param fov: Field of view radius in degrees
        :type fov: float
        :param freq: Observing frequency
        :type freq: float
        :param filter: Convolution filter to apply
        :type filter: `~nenupytv.image.AAFilter`
        :param cellsize: Cell size in degrees
        :type cellsize: float
    """

    def __init__(
        self,
        vis,
        uvw,
        freq,
        fov,
        conv_filter=AAFilter(),
        cellsize=None,
        robust=0.
        ):
        self.nsize = None
        self.fov = fov
        self.vis = vis
        self.uvw = uvw
        self.freq = freq
        self.filter = conv_filter
        self.cellsize = cellsize
        self.robust = robust

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def _ovsmpl(self):
        """ Filter oversampling
        """
        return self.filter.oversample

    @property
    def _hsup(self):
        """ Filter Half_sup
        """
        return self.filter.half_sup

    @property
    def _filter_idx(self):
        """ Filter indices
        """
        return np.arange(-self._hsup, self._hsup + 1)

    @property
    def _ftaps(self):
        """
        """
        return self.filter.filter_taps

    @property
    def vis(self):
        return self._vis
    @vis.setter
    def vis(self, v):
        x, y = np.tril_indices(v.shape[0], 0)
        xx = np.hstack((x[x!=y], y[x!=y]))
        yy = np.hstack((y[x!=y], x[x!=y]))
        self._vis = v[xx, yy, ...] + v[yy, xx, ...].conj()
        return

    @property
    def uvw(self):
        return self._uvw
    @uvw.setter
    def uvw(self, u):
        x, y = np.tril_indices(u.shape[0], 0)
        xx = np.hstack((x[x!=y], y[x!=y]))
        yy = np.hstack((y[x!=y], x[x!=y]))
        self._uvw = u[xx, yy, ...]

        maxu = np.max(np.abs(self._uvw[..., 0]))
        maxv = np.max(np.abs(self._uvw[..., 1]))
        self.resol = 1. / (5 * 2 * np.max((maxu, maxv))) * un.rad
        # self.resol = 1. / (10 * 2 * np.max((maxu, maxv))) * un.rad
        resol = self.resol.to(un.deg).value
        self.nsize = int(np.round(self.fov / resol))
        return

    @property
    def nsize(self):
        return self._nsize
    @nsize.setter
    def nsize(self, n):
        if n is None:
            self._nsize = None
        else:
            self._nsize = n
            self.measurement = np.zeros(
                (self.vis.shape[1], self.nsize, self.nsize),
                dtype='complex64'
            )
            self._meas_w = np.zeros(
                (self.nsize, self.nsize),
                dtype='float'
            )
            # for deconvolution the PSF should be 2x size of the image (see
            # Hogbom CLEAN for details), one grid for the sampling function:
            self.sampling = np.zeros(
                (2*self.nsize, 2*self.nsize),
                dtype='complex64'
            )
            self._samp_w = np.zeros(
                (2*self.nsize, 2*self.nsize),
                dtype='float'
            )

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    def populate(self):
        """
        """
        # if self.uvw.shape[0] != 1:
        #     raise ValueError(
        #         'UVW need to be averaged in time'
        #     )
        # if self.uvw.shape[1] != 1:
        #     raise ValueError(
        #         'UVW need to be averaged in frequency'
        #     )
        # iter over number of baselines in the triangular inferior
        # matrix (i.e. not autocorr), divided by 3 because (u, v, w)
        for vis_bl in range(int(self.uvw.size / 3)):
            u = self.uvw[vis_bl, 0] * np.radians(self.fov)
            v = self.uvw[vis_bl, 1] * np.radians(self.fov)

            if np.ma.is_masked(u) or np.ma.is_masked(v):
                continue

            du = int(np.round(u))
            dv = int(np.round(v))
            fu_offset = int(
                    (1 + self._hsup + (-u + du)) * self._ovsmpl
                )
            fv_offset = int(
                    (1 + self._hsup + (-v + dv)) * self._ovsmpl
                )

            du_psf = int(np.round(u*2))
            dv_psf = int(np.round(v*2))
            fu_offset_psf = int(
                    (1 + self._hsup + (-u*2 + du_psf)) * self._ovsmpl
                )
            fv_offset_psf = int(
                    (1 + self._hsup + (-v*2 + dv_psf)) * self._ovsmpl
            )

            if (dv + self.nsize // 2 + self._hsup >= self.nsize or
                du + self.nsize // 2 + self._hsup >= self.nsize or
                dv + self.nsize // 2 - self._hsup < 0 or
                du + self.nsize // 2 - self._hsup < 0):
                continue

            for conv_v in self._filter_idx:
                v_tap = self._ftaps[conv_v * self._ovsmpl + fv_offset]
                v_tap_psf = self._ftaps[conv_v * self._ovsmpl + fv_offset_psf]

                grid_v = dv + conv_v + self.nsize // 2
                grid_v_psf = dv_psf + conv_v + self.nsize

                for conv_u in self._filter_idx:
                    u_tap = self._ftaps[conv_u * self._ovsmpl + fu_offset]
                    u_tap_psf = self._ftaps[conv_u * self._ovsmpl + fu_offset_psf]

                    grid_u = du + conv_u + self.nsize // 2
                    grid_u_psf = du_psf + conv_u + self.nsize

                    conv_weight = v_tap * u_tap
                    conv_weight_psf = v_tap_psf * u_tap_psf

                    # for p in range(self.vis.shape[3]):
                    #     self.measurement[p, grid_v, grid_u] += self.vis[0, 0, vis_bl, p] * conv_weight
                    for p in range(self.vis.shape[1]):
                        self.measurement[p, grid_v, grid_u] += self.vis[vis_bl, p] * conv_weight
                        # self._meas_w[grid_v, grid_u] += 1
                        self._meas_w[grid_v, grid_u] += 1. * conv_weight
                    # assuming the PSF is the same for different correlations:
                    self.sampling[grid_v_psf, grid_u_psf] += (1+0.0j) * conv_weight_psf
                    # self._samp_w[grid_v_psf, grid_u_psf] += 1
                    self._samp_w[grid_v_psf, grid_u_psf] += 1. * conv_weight_psf

        self._compute_weights()
        return

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _compute_weights(self):
        """
        """
        # nbsl = self.uvw.size / 3
        # bsl_time = (2.*nbsl*1) # one time step
        # num = (5.*10.**(-self.robust))**2.*bsl_time
        # meas_f = num/np.sum(self._meas_w**2.)
        # samp_f = num/np.sum(self._samp_w**2.)

        # self.meas_weighted = self._meas_w / (1 + self._meas_w * meas_f)
        # self.samp_weighted = self._samp_w / (1 + self._samp_w * samp_f)

        factor = (5. * 10.**(-self.robust) )**2
        f = factor / (np.sum(self._meas_w**2.) / np.sum(self._meas_w))
        self.meas_weights = self._meas_w / (1 + self._meas_w * f)
        self.meas_weights /= self.meas_weights.max()

        f = factor / (np.sum(self._samp_w**2.) / np.sum(self._samp_w))
        self.samp_weights = self._samp_w / (1 + self._samp_w * f)
        self.samp_weights /= self.samp_weights.max()
        return
# ============================================================= #

# ============================================================= #
# ---------------------------- Grid --------------------------- #
# ============================================================= #
class Grid_Simple(object):
    """
    """

    def __init__(
        self,
        vis,
        uvw,
        freq,
        fov,
        cellsize=None,
        robust=0.,
        convolution=None
        ):
        self.nsize = None
        self.fov = fov
        self.vis = vis
        self.uvw = uvw
        self.freq = freq
        self.cellsize = cellsize
        self.robust = robust
        self.convolution = convolution

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def vis(self):
        return self._vis
    @vis.setter
    def vis(self, v):
        x, y = np.tril_indices(v.shape[0], 0)
        xx = np.hstack((x[x!=y], y[x!=y]))
        yy = np.hstack((y[x!=y], x[x!=y]))
        self._vis = v[xx, yy, ...] + v[yy, xx, ...].conj()
        return

    @property
    def uvw(self):
        return self._uvw
    @uvw.setter
    def uvw(self, u):
        x, y = np.tril_indices(u.shape[0], 0)
        xx = np.hstack((x[x!=y], y[x!=y]))
        yy = np.hstack((y[x!=y], x[x!=y]))
        self._uvw = u[xx, yy, ...]

        maxu = np.max(np.abs(self._uvw[..., 0]))
        maxv = np.max(np.abs(self._uvw[..., 1]))
        self.resol = 1. / (5 * 2 * np.max((maxu, maxv))) * un.rad
        resol = self.resol.to(un.deg).value
        self.nsize = int(np.round(self.fov / resol))
        return

    @property
    def nsize(self):
        return self._nsize
    @nsize.setter
    def nsize(self, n):
        if n is None:
            self._nsize = None
        else:
            self._nsize = n
            self.measurement = np.zeros(
                (self.vis.shape[1], self.nsize, self.nsize),
                dtype='complex64'
            )
            self._meas_w = np.zeros(
                (self.nsize, self.nsize),
                dtype='float'
            )
            # for deconvolution the PSF should be 2x size of the image (see
            # Hogbom CLEAN for details), one grid for the sampling function:
            self.sampling = np.zeros(
                (2*self.nsize, 2*self.nsize),
                dtype='complex64'
            )
            self._samp_w = np.zeros(
                (2*self.nsize, 2*self.nsize),
                dtype='float'
            )

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    def populate(self):
        """
        """
        for vis_bl in range(int(self.uvw.size / 3)):
            u = self.uvw[vis_bl, 0] * np.radians(self.fov)
            v = self.uvw[vis_bl, 1] * np.radians(self.fov)

            if np.ma.is_masked(u) or np.ma.is_masked(v):
                continue

            du = int(np.round(u)) + self.nsize // 2
            dv = int(np.round(v)) + self.nsize // 2
            du_psf = int(np.round(u*2)) + self.nsize
            dv_psf = int(np.round(v*2)) + self.nsize

            for p in range(self.vis.shape[1]):
                self.measurement[p, dv, du] += self.vis[vis_bl, p]
            self._meas_w[dv, du] += 1.
            # assuming the PSF is the same for different correlations:
            self.sampling[dv_psf, du_psf] += (1+0.0j)
            self._samp_w[dv_psf, du_psf] += 1.

        self._convolve()

        self._compute_weights()
        return

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _compute_weights(self):
        """
        """
        factor = (5. * 10.**(-self.robust) )**2
        f = factor / (np.sum(self._meas_w**2.) / np.sum(self._meas_w))
        self.meas_weights = self._meas_w / (1 + self._meas_w * f)
        self.meas_weights /= self.meas_weights.max()

        f = factor / (np.sum(self._samp_w**2.) / np.sum(self._samp_w))
        self.samp_weights = self._samp_w / (1 + self._samp_w * f)
        self.samp_weights /= self.samp_weights.max()
        return

    def _convolve(self):
        """
        """
        if self.convolution is None:
            return
        elif self.convolution.lower() == 'gaussian':
            from scipy.ndimage import gaussian_filter
            def _scipy_gauss(im, sig=0.25):
                if im.dtype == np.complex64:
                    im_re = gaussian_filter(
                        input=im.real,
                        sigma=sig,
                        order=0,
                        output=None,
                        mode='reflect',
                        cval=0.0,
                        truncate=4.0
                    )
                    im_im = gaussian_filter(
                        input=im.imag,
                        sigma=sig,
                        order=0,
                        output=None,
                        mode='reflect',
                        cval=0.0,
                        truncate=4.0
                    )
                    return im_re + 1.j*im_im
                else:
                    return gaussian_filter(
                        input=im,
                        sigma=sig,
                        order=0,
                        output=None,
                        mode='reflect',
                        cval=0.0,
                        truncate=4.0
                    )
            for p in range(self.measurement.shape[0]):
                self.measurement[p, ...] = _scipy_gauss(self.measurement[p, ...])
            self.sampling = _scipy_gauss(self.sampling)
            self._meas_w = _scipy_gauss(self._meas_w)
            self._samp_w = _scipy_gauss(self._samp_w)
        else:
            raise Exception(
                'Convolution kernel {} not implemented.'.format(self.convolution)
            )
# ============================================================= #

# ============================================================= #
# -------------------------- Grid_HPX ------------------------- #
# ============================================================= #
class Grid_HPX(object):
    """
    """

    def __init__(self, resolution):
        self.nside = None
        self.resol = resolution
        self.ra, self.dec = self._hpx_radec()

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def resol(self):
        """ Resolution of the cell in degrees
        """
        return self._resol
    @resol.setter
    def resol(self, r):
        nsides = 2**np.arange(1, 12)
        resol_rad = hp.nside2resol(
            nside=nsides,
            arcmin=False
            )
        resol_deg = np.degrees(resol_rad)
        idx = (np.abs(resol_deg - r)).argmin()
        self._resol = resol_deg[idx]
        self.nside = nsides[idx]
        return

    @property
    def nside(self):
        return self._nside
    @nside.setter
    def nside(self, n):
        if n is None:
            self._nside = None
        else:
            self.npix = hp.nside2npix(n)
            self._nside = n
        return

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _hpx_radec(self):
        """
        """
        ra, dec = hp.pix2ang(
            nside=self.nside,
            ipix=np.arange(self.npix),
            lonlat=True,
            nest=False
            )
        return ra, dec

# ============================================================= #

# class ImgWPlus(aipy.img.ImgW):
#     """
#     Sub-class of the aipy.img.ImgW class that adds support for different
#     visibility weighting scheme and uv plane tapering.  This class also
#     adds in a couple of additional methods that help determine the size of
#     the field of view and the pixels near the phase center.
#     """
#     def __init__(self, size=100, res=1, wres=.5, mf_order=0):
#         """size = number of wavelengths which the UV matrix spans (this
#         determines the image resolution).
#         res = resolution of the UV matrix (determines image field of view).
#         wres: the gridding resolution of sqrt(w) when projecting to w=0."""
#         self.res = float(res)
#         self.size = float(size)
#         ## Small change needed to work with Numpy 1.12+
#         dim = numpy.int64(numpy.round(self.size / self.res))
#         self.shape = (dim,dim)
#         self.uv = numpy.zeros(shape=self.shape, dtype=numpy.complex64)
#         self.bm = []
#         for i in range(mf_order+1):
#             self.bm.append(numpy.zeros(shape=self.shape, dtype=numpy.complex64))
#         self.wres = float(wres)
#         self.wcache = {}
#     def put(self, uvw, data, wgts=None, invker2=None):
#         """Same as Img.put, only now the w component is projected to the w=0
#         plane before applying the data to the UV matrix."""
#         u, v, w = uvw
#         if len(u) == 0: return
#         if wgts is None:
#             wgts = []
#             for i in range(len(self.bm)):
#                 if i == 0:
#                     wgts.append(numpy.ones_like(data))
#                 else:
#                     wgts.append(numpy.zeros_like(data))
#         if len(self.bm) == 1 and len(wgts) != 1:
#             wgts = [wgts]
#         assert(len(wgts) == len(self.bm))
#         # Sort uvw in order of w
#         order = numpy.argsort(w)
#         u = u.take(order)
#         v = v.take(order)
#         w = w.take(order)
#         data = data.take(order)
#         wgts = [wgt.take(order) for wgt in wgts]
#         sqrt_w = numpy.sqrt(numpy.abs(w)) * numpy.sign(w)
#         i = 0
#         while True:
#             # Grab a chunk of uvw's that grid w to same point.
#             j = sqrt_w.searchsorted(sqrt_w[i]+self.wres)
#             print('%d/%d datums' % (j, len(w)))
#             avg_w = numpy.average(w[i:j])
#             # Put all uv's down on plane for this gridded w point
#             wgtsij = [wgt[i:j] for wgt in wgts]
#             uv,bm = aipy.img.Img.put(self, (u[i:j],v[i:j],w[i:j]),
#                 data[i:j], wgtsij, apply=False)
#             # Convolve with the W projection kernel
#             invker = numpy.fromfunction(lambda u,v: self.conv_invker(u,v,avg_w),
#                 uv.shape)
#             if not invker2 is None:
#                 invker *= invker2
#             self.uv += ifft2Function(fft2Function(uv) * invker)
#             #self.uv += uv
#             for b in range(len(self.bm)):
#                 self.bm[b] += ifft2Function(fft2Function(bm[b]) * invker)
#                 #self.bm[b] += numpy.array(bm)[0,:,:]
#             if j >= len(w):
#                 break
#             i = j


#!/usr/bin/env python3
import sys
import time
import numpy.linalg
import numpy.random
import random

VK_SPREAD = [
    [-24, 58.0], [-23, 36.0], [-22, 20.0], [-21, 17.0],
    [-20, 16.0], [-19, 19.0], [-18, 24.0], [-17, 36.0],
    [-16, 45.0], [-15, 55.0], [-14, 60.0], [-13, 70.0],
    [-12, 75.0], [-11, 79.0], [-10, 83.0], [-9, 87.0],
    [-8, 92.0], [-7, 96.0], [-6, 100.0], [-5, 105.0],
    [-4, 108.0], [-3, 105.0], [-2, 96.0], [-1, 80.0],
    [0, 58.0], [1, 36.0], [2, 20.0], [3, 17.0],
    [4, 16.0], [5, 19.0], [6, 24.0], [7, 36.0],
    [8, 45.0], [9, 55.0], [10, 60.0], [11, 70.0],
    [12, 75.0], [13, 79.0], [14, 83.0], [15, 87.0],
    [16, 92.0], [17, 96.0], [18, 100.0], [19, 105.0],
    [20, 108.0], [21, 105.0], [22, 96.0], [23, 80.0],
    [24, 58.0], [25, 36.0], [26, 20.0], [27, 17.0],
    [28, 16.0], [29, 19.0], [30, 24.0], [31, 36.0],
    [32, 45.0], [33, 55.0], [34, 60.0], [35, 70.0],
    [36, 75.0], [37, 79.0], [38, 83.0], [39, 87.0],
    [40, 92.0], [41, 96.0], [42, 100.0], [43, 105.0],
    [44, 108.0], [45, 105.0], [46, 96.0], [47, 80.0],
    [48, 58.0],
]

def WR_params(date, f=1):
    """
        Generates params for WeightRandom class
    """
    tm = time.localtime(date)
    date = (tm.tm_hour * 60 + tm.tm_min) * 60 + tm.tm_sec
    return [
        [date, 0.0],
        [date + (3600 + 600) / f, 3200],
        [date + (2 * 3600) / f, 1850]
    ]

class PyWeightRandom:
    """
        use post date and time of publish
        to generate weight random method()
        it emulates when user creates reposts/likes
    """

    def __init__(self, data):
        """
            data = [
              [ x0 , y0 ] , - левая точка нуля пораболы
              [ x1 , y1 ] , - пик
              [ x2 , y2 ] - точка после точки пика через ( x1-x0 )
            ]
            endpoint = mun of second that:
             1. > x0
             2. <= x0 + 86000
        """
        # p = [ x0 , x1 , y1 , y2]
        self.p = [data[0][0], data[1][0], data[1][1], data[2][1]]
        mm = [
            [self.p[0]**2, self.p[0], 1],
            [self.p[1]**2, self.p[1], 1],
            [(2 * self.p[1] - self.p[0])**2, (2 * self.p[1] - self.p[0]), 1]
        ]
        abc = list(numpy.linalg.solve(mm, [0, self.p[2], 0]))
        self.a = abc[0]
        self.b = abc[1]
        self.c = abc[2]
        mm = [
            [1.0, -self.p[2]],
            [1.0, -self.p[3]]
        ]
        dk = numpy.linalg.solve(mm, [self.p[2] * self.p[1], self.p[3] * (2 * self.p[1] - self.p[0])])
        self.d = dk[0]
        self.k = dk[1]

        self.data = {}
        self._fill()

        self._weight = list(self.data.values())
        self._m = sum(self._weight)
        self.elem = list(self.data.keys())
        self.weight = list(map(lambda x: float(x) / float(self._m), self._weight))

    # ==========================================================================
    #                            INTERNAL METHODS
    # ==========================================================================

    def spread(self, x):
        """
            формула распределения случайно величиный
            p = [
             0 - start point (flaot)
             1 - X where max of Y (float)
             2 - max of Y (float)
             3 - Y when x = ( 2*p[1]-p[0] )
            ]
        """
        if self.p[0] + 1 <= x <= self.p[1] + 1:
            res = self.a * (x**2) + self.b * x + self.c
        elif self.p[1] <= x:
            res = (self.d / (x + self.k))
        else:
            res = self.spread(x + 3600 * 24)
        return res

    def _fill(self):
        """
            should fill self.data with values of spread function in interesing range
            return None
        """
        end = 86400 + (i := self.p[0])
        while i < end:
            self.data[i % 86400] = self.spread(i)
            i += 1

    def random(self, y=1):
        """
            return random number using spread function
        """
        if y <= 1:
            x = random.random()
            i = 0
            while x > 0.0:
                x -= self.weight[i]
                i += 1
            return [i - 1]
        else:
            k = float(self._m) / float(y)
            x = []
            y = k
            for i in range(len(ww := list(self._weight))):
                while ww[i] > 0.0:
                    y -= (z := min(y, ww[i]))
                    ww[i] -= z
                    if y <= 0.0:
                        y = k
                        x.append(ww[i])
            return x

    def destruct(self):
        """
            really do nothing (here)
        """
        pass

    def will_be(self, count_now, now):
        """
            predict now many obj will be in 24 hours
        """
        def en(_time):
            tm = time.localtime(_time)
            return (tm.tm_hour * 60 + tm.tm_min) * 60 + tm.tm_sec
        if (now := en(now)) < self.p[0]:
            now += 24 * 3600

        return (
            (count_now * sum(self.weight) / s)
            if (
                s := sum(
                    self.weight[i]
                    for i in range(len(self.weight))
                    if self.p[0] < i < now
                )
            ) > 0 else
            0
        )

def py__interpolice(ff, x):
    """
        real interpolice method
        USE ONLY BY INTERPOLICE
    """
    res = 0.0
    for i in ff:
        if not(-12 < (i[0] - x) < 12):
            continue
        li = 1.0
        for j in ff:
            if not(-12 < (j[0] - x) < 12):
                continue
            if i[0] != j[0]:
                z = (i[0] - j[0])
                li *= ((x - j[0]) / z)
        res += li * i[1]
    return res

boo = False
if sys.platform in {'linux', 'darwin'} and '--no-c' not in sys.argv[1:]:
    try:
        from .ext import interpolice as __interpolice
        from .ext import CWeightRandom as WeightRandom
        boo = True
    except Exception:
        pass

if not boo:
    WeightRandom = PyWeightRandom
    __interpolice = py__interpolice

def interpolice(x, ff=None):
    """
        takes as x float() - number of hours
          13:30 - 13.5
        return float from [0..1]
        ff = [
         [x0 , f(x0)],
         [x1 , f(x1)],
        ]
    """
    if ff is None:
        ff = VK_SPREAD
        _max = 108.0
    else:
        _max = float(max(i[1] for i in ff))
    return __interpolice(ff, x) / _max


# -*- coding: utf-8 -*-
# ---------------------

import numpy as np
import ergonomics.utils as utils

class RebaScore:
    '''
    Class to compute REBA metrics

    Pose:
          [0]: Head
          [1]: Neck
          [2, 3, 4, 14]: Left arm + (optional)left hand
          [5, 6, 7, 15]: Right arm + (optional)right hand
          [8, 9, 10]: Left leg
          [11, 12, 13]: Right leg
    '''
    def __init__(self):
        # Table A ( Neck X Trunk X Legs)
        self.table_a = np.zeros((3, 5, 4))
        # Table B ( UpperArm X LowerArm X Wrist)
        self.table_b = np.zeros((6, 2, 3))
        # Table C ( ScoreA X ScoreB)
        self.table_c = np.zeros((12, 12))

        # Body Params
        self.body = {'neck_angle': 0, 'neck_side': False,
                     'trunk_angle': 0, 'trunk_side': False,
                     'legs_walking': False, 'legs_angle': 0,
                     'load': 0}

        # Arms Params
        self.arms = {'upper_arm_angle': 0, 'shoulder_raised': False, 'arm_abducted': False, 'leaning': False,
                     'lower_arm_angle': 0,
                     'wrist_angle': 0, 'wrist_twisted': False}

        # Init lookup tables
        self.init_table_a()
        self.init_table_b()
        self.init_table_c()

    def init_table_a(self):
        '''
        Table used to compute upper body score

        :return: None
        '''
        self.table_a = np.array([
                                [[1, 2, 3, 4], [2, 3, 4, 5], [2, 4, 5, 6], [3, 5, 6, 7], [4, 6, 7, 8]],
                                [[1, 2, 3, 4], [3, 4, 5, 6], [4, 5, 6, 7], [5, 6, 7, 8], [6, 7, 8, 9]],
                                [[3, 3, 5, 6], [4, 5, 6, 7], [5, 6, 7, 8], [6, 7, 8, 9], [7, 8, 9, 9]]
                                ])

    def init_table_b(self):
        '''
        Table used to computer lower body score

        :return: None
        '''
        self.table_b = np.array([
                                [[1, 2, 2], [1, 2, 3]],
                                [[1, 2, 3], [2, 3, 4]],
                                [[3, 4, 5], [4, 5, 5]],
                                [[4, 5, 5], [5, 6, 7]],
                                [[6, 7, 8], [7, 8, 8]],
                                [[7, 8, 8], [8, 9, 9]],
                                ])

    def init_table_c(self):
        '''
        Table to compute score_c

        :return: None
        '''
        self.table_c = np.array([
                                [1, 1, 1, 2, 3, 3, 4, 5, 6, 7, 7, 7],
                                [1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 7, 8],
                                [2, 3, 3, 3, 4, 5, 6, 7, 7, 8, 8, 8],
                                [3, 4, 4, 4, 5, 6, 7, 8, 8, 9, 9, 9],
                                [4, 4, 4, 5, 6, 7, 8, 8, 9, 9, 9, 9],
                                [6, 6, 6, 7, 8, 8, 9, 9, 10, 10, 10, 10],
                                [7, 7, 7, 8, 9, 9, 9, 10, 10, 11, 11, 11],
                                [8, 8, 8, 9, 10, 10, 10, 10, 10, 11, 11, 11],
                                [9, 9, 9, 10, 10, 10, 11, 11, 11, 12, 12, 12],
                                [10, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12, 12],
                                [11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 12],
                                [12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12],
                                ])

    def set_body(self, values):
        # type: (np.ndarray) -> None
        '''
        Set body params

        :param values: [neck_angle, neck_side, trunk_angle, trunk_side,
                        legs_walking, legs_angle, load]

        :return: None
        '''
        assert len(values) == len(self.body)

        for i, (key, _) in enumerate(self.body.items()):
            self.body[key] = values[i]

    def set_arms(self, values):
        # type: (np.ndarray) -> None
        '''
        Set arms params

        :param values:  [upper_arm_angle, shoulder_raised, arm_abducted, leaning,
                        lower_arm_angle, wrist_angle, wrist_twisted]

        :return: None
        '''
        assert len(values) == len(self.arms)

        for i, (key, _) in enumerate(self.arms.items()):
            self.arms[key] = values[i]

    def compute_score_a(self):
        # type: (RebaScore) -> (np.ndarray, np.ndarray)
        '''
        Compute score A
        >>> rebascore = RebaScore()
        >>> rebascore.set_body(np.array([10, 0, 20, 0, 1, 50, 0]))
        >>> rebascore.compute_score_a()
        (4, array([1, 2, 3]))

        :return: Score A, [neck_score, trunk_score, leg_score]
        '''
        neck_score, trunk_score, leg_score, load_score = 0, 0, 0, 0

        # Neck position
        if 10 <= self.body['neck_angle'] <= 20 :
            neck_score +=1
        else:
            neck_score +=2
        # Neck adjust
        neck_score +=1 if self.body['neck_side'] else 0

        # Trunk position
        if 0 <= self.body['trunk_angle'] <= 1:
            trunk_score +=1
        elif self.body['trunk_angle'] <= 20:
            trunk_score +=2
        elif 20 <= self.body['trunk_angle'] <= 60:
            trunk_score +=3
        elif self.body['trunk_angle'] > 60:
            trunk_score +=4
        # Trunk adjust
        trunk_score += 1 if self.body['trunk_side'] else 0

        # Legs position
        leg_score += 2 if self.body['legs_walking'] else 1
        # Legs adjust
        if 30 <= self.body['legs_angle'] <= 60:
            leg_score += 1
        elif self.body['legs_angle'] > 60:
            leg_score += 2

        # Load
        if 5 <= self.body['load'] <= 10:
            load_score += 1
        elif self.body['load'] > 10:
            load_score += 2

        assert neck_score > 0 and trunk_score > 0 and leg_score > 0

        score_a = self.table_a[neck_score-1][trunk_score-1][leg_score-1]
        return score_a, np.array([neck_score, trunk_score, leg_score])

    def compute_score_b(self):
        # type: (RebaScore) -> (np.ndarray, np.ndarray)
        '''
        Compute score B
        >>> rebascore = RebaScore()
        >>> rebascore.set_arms(np.array([45, 0, 0, 0, 70, 0, 1]))
        >>> rebascore.compute_score_b()
        (2, array([2, 1, 2]))

        :return: scoreB, [upper_arm_score, lower_arm_score, wrist_score]
        '''
        upper_arm_score, lower_arm_score, wrist_score = 0, 0, 0

        # Upper arm position
        if -20 <= self.arms['upper_arm_angle'] <= 20:
            upper_arm_score +=1
        elif self.arms['upper_arm_angle'] <= 45:
            upper_arm_score +=2
        elif 45 <= self.arms['upper_arm_angle'] <= 90:
            upper_arm_score +=3
        elif self.arms['upper_arm_angle'] > 90:
            upper_arm_score +=4

        # Upper arm adjust
        upper_arm_score += 1 if self.arms['shoulder_raised'] else 0
        upper_arm_score += 1 if self.arms['arm_abducted'] else 0
        upper_arm_score -= 1 if self.arms['leaning'] else 0

        # Lower arm position
        if 60 <= self.arms['lower_arm_angle'] <= 100:
            lower_arm_score += 1
        else:
            lower_arm_score += 2

        # Wrist position
        if -15 <= self.arms['wrist_angle'] <= 15:
            wrist_score += 1
        else:
            wrist_score += 2

        # Wrist adjust
        wrist_score += 1 if self.arms['wrist_twisted'] else 0

        assert lower_arm_score > 0 and wrist_score > 0

        score_b = self.table_b[upper_arm_score-1][lower_arm_score-1][wrist_score-1]
        return score_b, np.array([upper_arm_score, lower_arm_score, wrist_score])

    def compute_score_c(self, score_a, score_b):
        # type: (np.ndarray, np.ndarray) -> (np.ndarray, str)
        '''
        Compute score C

        :param score_a:  Score A
        :param score_b:  Score B

        :return: Score C, caption
        '''
        reba_scoring = ['Negligible Risk',
                         'Low Risk. Change may be needed',
                         'Medium Risk. Further Investigate. Change Soon',
                         'High Risk. Investigate and Implement Change',
                         'Very High Risk. Implement Change'
                         ]

        score_c = self.table_c[score_a-1][score_b-1]
        ix = self.score_c_to_5_classes(score_c)
        caption = reba_scoring[ix]

        return score_c, caption

    @staticmethod
    def score_c_to_5_classes(score_c):
        # type: (np.ndarray) -> int
        '''
        Score C to 5 risk-classes

        :param score_c:  Score C
        :return: Risk-class
        '''
        if score_c == 1:
            ret = 0
        elif 2 <= score_c <= 3:
            ret = 1
        elif 4 <= score_c <= 7:
            ret = 2
        elif 8 <= score_c <= 10:
            ret = 3
        else:
            ret = 4

        return ret

    @staticmethod
    def get_body_angles_from_pose_left(pose, verbose=False):
        # type: (np.ndarray, bool) -> np.ndarray
        '''
        Get body angles from pose (look at left)

        :param pose: Pose (Joints coordinates)
        :param verbose: If true show each pose for debugging

        :return: Body params (neck_angle, neck_side, trunk_angle, trunk_side,
                legs_walking, legs_angle, load)
        '''

        pose = np.expand_dims(np.copy(pose), 0)

        neck_angle, neck_side, trunk_angle, trunk_side, \
        legs_walking, legs_angle, load = 0, 0, 0, 0, 0, 0, 0

        if verbose:
            utils.show_skeleton(pose, title="GT pose")
            _pose, _ = utils.rotate_pose(np.copy(pose), rotation_joint=8)
            utils.show_skeleton(_pose, title="GT pose left")

        # Trunk position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8)
        pose -= (pose[:, 8] + pose[:, 11]) /2

        if quad(pose[0, 1]) < 3:
            trunk_angle = np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]) - (np.pi / 2))
        else:
            trunk_angle = 270 + np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]))

        if verbose:
            utils.show_skeleton(pose, title="Trunk angle: " + str(round(trunk_angle, 2)))

        # Trunk bending
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi/2)

        angle = np.pi /2 if quad(pose[0, 1]) > 2 else -np.pi/2
        trunk_side_angle = abs(np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]) + angle))
        trunk_side = 1 if trunk_side_angle > 30 else 0

        if verbose:
            utils.show_skeleton(pose, title="Trunk side angle: " + str(round(trunk_side_angle, 2)))

        # Neck position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi/2)
        pose -= pose[:, 1]

        if quad(pose[0, 0]) < 3:
            neck_angle = np.rad2deg(np.arctan2(pose[0, 0, 1], pose[0, 0, 0]) - (np.pi / 2)) - trunk_angle
        else:
            neck_angle = 270 + np.rad2deg(np.arctan2(pose[0, 0, 1], pose[0, 0, 0])) - trunk_angle

        if verbose:
            utils.show_skeleton(pose, title="Neck angle: " + str(round(neck_angle, 2)))

        # Neck bending
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi / 2)
        angle = np.pi /2 if quad(pose[0, 0]) > 2 else -np.pi/2
        neck_side_angle = abs(np.rad2deg(np.arctan2(pose[0, 0, 1], pose[0, 0, 0]) + angle)) - trunk_side_angle
        neck_side = 1 if neck_side_angle > 20 else 0

        if verbose:
            utils.show_skeleton(pose, title="Neck side angle: " + str(round(neck_side_angle, 2)))

        # Legs position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
        pose -= pose[:, 8]

        if quad(pose[0, 9]) > 2:
            legs_angle = -np.rad2deg(np.arctan2(pose[0, 9, 1], pose[0, 9, 0]) + (np.pi/2))
        else:
            legs_angle = 270 - np.rad2deg(np.arctan2(pose[0, 9, 1], pose[0, 9, 0]))

        step_size = abs(np.linalg.norm(pose[0, 10, :2] - pose[0, 13, :2]))
        legs_walking = 1 if step_size > 0.1 else 0

        if verbose:
            title = "Leg angle: " + str(round(legs_angle, 2)) + " Step size: " + str(round(step_size, 2))
            utils.show_skeleton(pose, title=title)

        return np.array([neck_angle, neck_side, trunk_angle, trunk_side,
                           legs_walking, legs_angle, load])

    @staticmethod
    def get_body_angles_from_pose_right(pose, verbose=False):
        # type: (np.ndarray, bool) -> np.ndarray
        '''
        Get body angles from pose (look at right)

        :param pose: Pose (Joints coordinates)
        :param verbose: If true show each pose for debugging

        :return: Body params (neck_angle, neck_side, trunk_angle, trunk_side,
                              legs_walking, legs_angle, load)
        '''
        pose = np.expand_dims(np.copy(pose), 0)

        neck_angle, neck_side, trunk_angle, trunk_side, \
        legs_walking, legs_angle, load = 0, 0, 0, 0, 0, 0, 0

        if verbose:
            utils.show_skeleton(pose, title="GT pose")
            _pose, _ = utils.rotate_pose(np.copy(pose), rotation_joint=8)
            _pose, _ = utils.rotate_pose(_pose, rotation_joint=8, m_coeff=np.pi)
            utils.show_skeleton(_pose, title="GT pose Right")

        pose, _ = utils.rotate_pose(pose, rotation_joint=8)

        # Trunk position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi)
        pose -= (pose[:, 8] + pose[:, 11]) / 2
        trunk_angle = np.rad2deg((np.pi / 2)  - np.arctan2(pose[0, 1, 1], pose[0, 1, 0]))

        if verbose:
            utils.show_skeleton(pose, title="Trunk angle: " + str(round(trunk_angle, 2)))

        # Trunk bending
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi / 2)
        trunk_side_angle = abs(np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]) - (np.pi / 2)))
        trunk_side = 1 if trunk_side_angle > 30 else 0

        if verbose:
            utils.show_skeleton(pose, title="Trunk side angle: " + str(round(trunk_side_angle, 2)))

        # Neck position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
        pose -= pose[:, 1]
        neck_angle = np.rad2deg((np.pi / 2) - np.arctan2(pose[0, 0, 1], pose[0, 0, 0])) - trunk_angle

        if verbose:
            utils.show_skeleton(pose, title="Neck angle: " + str(round(neck_angle, 2)))

        # Neck bending
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi / 2)
        neck_side_angle = np.abs(np.rad2deg(np.abs(np.arctan2(pose[0, 0, 1], pose[0, 0, 0])) - (np.pi / 2)))
        neck_side = 1 if neck_side_angle > 20 else 0

        if verbose:
            utils.show_skeleton(pose, title="Neck side angle: " + str(round(neck_side_angle, 2)))

        # Legs position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
        pose -= pose[:, 11]
        legs_angle = np.rad2deg((np.pi / 2) + np.arctan2(pose[0, 12, 1], pose[0, 12, 0]))
        step_size = abs(np.linalg.norm(pose[0, 10, :2] - pose[0, 13, :2]))
        legs_walking = 1 if step_size > 0.1 else 0

        if verbose:
            title = "Leg angle: " + str(round(legs_angle, 2)) + " Step size: " + str(round(step_size, 2))
            utils.show_skeleton(pose, title=title)

        return np.array([neck_angle, neck_side, trunk_angle, trunk_side,
                legs_walking, legs_angle, load])

    @staticmethod
    def get_arms_angles_from_pose_left(pose, verbose=False):
        # type: (np.ndarray, bool) -> np.ndarray
        '''
        Get arms angles from pose (look at left)

        :param pose: Pose (Joints coordinates)
        :param verbose: If true show each pose for debugging

        :return: Body params (upper_arm_angle, shoulder_raised, arm_abducted, leaning,
                              lower_arm_angle, wrist_angle, wrist_twisted)
        '''
        pose = np.expand_dims(np.copy(pose), 0)
        if verbose:
            utils.show_skeleton(pose, title="GT pose")

        # Leaning
        pose, _ = utils.rotate_pose(pose, rotation_joint=8)
        pose -= (pose[:, 8] + pose[:, 11]) / 2

        if quad(pose[0, 1]) < 3:
            trunk_angle = np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]) - (np.pi / 2))
        else:
            trunk_angle = 270 + np.rad2deg(np.arctan2(pose[0, 1, 1], pose[0, 1, 0]))

        leaning = 1 if trunk_angle > 30 else 0

        if verbose:
            utils.show_skeleton(pose, title="Leaning angle: " + str(round(trunk_angle, 2)))

        # Upper Arm position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8)
        pose -= pose[:, 2]

        if quad(pose[0, 3]) > 2:
            upper_arm_angle = -np.rad2deg(np.arctan2(pose[0, 3, 1], pose[0, 3, 0]) + (np.pi / 2))
        else:
            upper_arm_angle = 270 - np.rad2deg(np.arctan2(pose[0, 3, 1], pose[0, 3, 0]))

        upper_arm_angle += trunk_angle

        if verbose:
            utils.show_skeleton(pose, title="Upper Arms angle: " + str(round(upper_arm_angle, 2)))

        # Upper Arm Adjust
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi/2)
        shoulder_step = pose[:, 2, 1] - pose[:, 1, 1]

        if quad(pose[0, 3]) > 2:
            arm_abducted_angle = -np.rad2deg(np.arctan2(pose[0, 3, 1], pose[0, 3, 0]) + (np.pi / 2))
        else:
            arm_abducted_angle = 270 - np.rad2deg(np.arctan2(pose[0, 3, 1], pose[0, 3, 0]))

        shoulder_raised = 1 if shoulder_step > 0.02 else 0
        arm_abducted = 1 if arm_abducted_angle > 45 else 0

        if verbose:
            print(shoulder_raised)
            utils.show_skeleton(pose, title="Upper Arms abducted: " + str(round(arm_abducted_angle, 2)))

        # Lower Arm position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi/2)
        pose -= pose[:, 3]

        if quad(pose[0, 4]) > 2:
            lower_arm_angle = -np.rad2deg(np.arctan2(pose[0, 4, 1], pose[0, 4, 0]) + (np.pi / 2))
        else:
            lower_arm_angle = 270 - np.rad2deg(np.arctan2(pose[0, 4, 1], pose[0, 4, 0]))

        lower_arm_angle = lower_arm_angle + trunk_angle - upper_arm_angle

        if verbose:
            utils.show_skeleton(pose, title="Lower Arms angle: " + str(round(lower_arm_angle, 2)))

        # Wrist position
        wrist_angle = 0
        wrist_twisted = 0

        if pose.shape[1] > 14:
            pose -= pose[:, 4]
            wrist_angle = np.rad2deg(np.arctan2(pose[0, 14, 1], pose[0, 14, 0]) - (np.pi / 2) )

            if verbose:
                utils.show_skeleton(pose, title="Wrist Angle: " + str(round(wrist_angle, 2)))

            pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
            wrist_twisted_angle = abs(np.rad2deg(np.arctan2(pose[0, 14, 1], pose[0, 14, 0]) - (np.pi / 2)))
            wrist_twisted = 1 if wrist_twisted_angle > 30 else 0

        return np.array([upper_arm_angle, shoulder_raised, arm_abducted, leaning,
                lower_arm_angle, wrist_angle, wrist_twisted])

    @staticmethod
    def get_arms_angles_from_pose_right(pose, verbose=False):
        # type: (np.ndarray, bool) -> np.ndarray
        '''
        Get arms angles from pose (look at right)

        :param pose: Pose (Joints coordinates)
        :param verbose: If true show each pose for debugging

        :return: Body params (upper_arm_angle, shoulder_raised, arm_abducted, leaning,
                               lower_arm_angle, wrist_angle, wrist_twisted)
        '''
        pose = np.expand_dims(np.copy(pose), 0)
        if verbose:
            utils.show_skeleton(pose, title="GT pose")

        pose, _ = utils.rotate_pose(pose, rotation_joint=8)
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi)

        # Leaning
        pose -= (pose[:, 8] + pose[:, 11]) / 2
        trunk_angle = np.rad2deg((np.pi / 2) - np.arctan2(pose[0, 1, 1], pose[0, 1, 0]))
        leaning = 1 if trunk_angle > 60 else 0

        if verbose:
            utils.show_skeleton(pose, title="Leaning angle: " + str(round(trunk_angle, 2)))

        # Upper Arm position
        pose -= pose[:, 5]
        if quad(pose[0, 6]) == 2:
            upper_arm_angle = -(270 - np.rad2deg(np.arctan2(pose[0, 6, 1], pose[0, 6, 0])) - trunk_angle)
        else:
            upper_arm_angle = np.rad2deg((np.pi /2) + np.arctan2(pose[0, 6, 1], pose[0, 6, 0])) + trunk_angle

        if verbose:
            utils.show_skeleton(pose, title="Upper Arms angle: " + str(round(upper_arm_angle, 2)))

        # Upper Arm Adjust
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=np.pi / 2)
        shoulder_step = pose[:, 5, 1] - pose[:, 1, 1]
        arm_abducted_angle = abs(np.rad2deg((np.pi / 2) + np.arctan2(pose[0, 6, 1], pose[0, 6, 0])))
        shoulder_raised = 1 if shoulder_step > 0.02 else 0
        arm_abducted = 1 if arm_abducted_angle > 45 else 0

        if verbose:
            print(shoulder_raised)
            utils.show_skeleton(pose, title="Upper Arms abducted: " + str(round(arm_abducted_angle, 2)))

        # Lower Arm position
        pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
        pose -= pose[:, 6]
        lower_arm_angle = np.rad2deg((np.pi / 2) + np.arctan2(pose[0, 7, 1], pose[0, 7, 0]) ) + trunk_angle - upper_arm_angle

        if verbose:
            utils.show_skeleton(pose, title="Lower Arms angle: " + str(round(lower_arm_angle, 2)))

        # Wrist position
        wrist_angle = 0
        wrist_twisted = 0

        if pose.shape[1] > 14:
            pose -= pose[:, 7]
            wrist_angle = np.rad2deg((np.pi / 2) + np.arctan2(pose[0, 15, 1], pose[0, 15, 0]))

            if verbose:
                utils.show_skeleton(pose, title="Wrist Angle: " + str(round(wrist_angle, 2)))

            pose, _ = utils.rotate_pose(pose, rotation_joint=8, m_coeff=-np.pi / 2)
            wrist_twisted_angle = abs(np.rad2deg((np.pi / 2) + np.arctan2(pose[0, 15, 1], pose[0, 15, 0])))
            wrist_twisted = 1 if wrist_twisted_angle > 30 else 0

        return np.array([upper_arm_angle, shoulder_raised, arm_abducted, leaning,
                         lower_arm_angle, wrist_angle, wrist_twisted])

def quad(coord):
    q = 0
    if coord[0] >= 0 and coord[1] >= 0:
        q = 1
    elif coord[0] <= 0 and coord[1] >= 0:
        q = 2
    elif coord[0] <= 0 and coord[1] <= 0:
        q = 3
    elif coord[0] >= 0 and coord[1] <= 0:
        q = 4
    return q

if __name__ == '__main__':

    import doctest
    doctest.testmod()

    sample_pose = np.array([[ 0.08533354,  1.03611605,  0.09013124],
                      [ 0.15391247,  0.91162637, -0.00353906],
                      [ 0.22379057,  0.87361878,  0.11541229],
                      [ 0.4084777 ,  0.69462843,  0.1775224 ],
                      [ 0.31665226,  0.46389668,  0.16556387],
                      [ 0.1239769 ,  0.82994377, -0.11715403],
                      [ 0.08302169,  0.58146328, -0.19830338],
                      [-0.06767788,  0.53928527, -0.00511249],
                      [ 0.11368726,  0.49372503,  0.21275574],
                      [ 0.069179  ,  0.07140968,  0.26841402],
                      [ 0.10831762, -0.36339359,  0.34032449],
                      [ 0.11368726,  0.41275504, -0.01171348],
                      [ 0.        ,  0.        ,  0.        ],
                      [ 0.02535541, -0.43954643,  0.04373671],
                      [ 0.26709431,  0.33643749,  0.17985192],
                      [-0.15117603,  0.49462711,  0.02703403]])

    rebaScore = RebaScore()

    body_params = rebaScore.get_body_angles_from_pose_right(sample_pose)
    arms_params = rebaScore.get_arms_angles_from_pose_right(sample_pose)

    rebaScore.set_body(body_params)
    score_a, partial_a = rebaScore.compute_score_a()

    rebaScore.set_arms(arms_params)
    score_b, partial_b = rebaScore.compute_score_b()

    score_c, caption = rebaScore.compute_score_c(score_a, score_b)

    print("Score A: ", score_a, "Partial: ", partial_a)
    print("Score A: ", score_b, "Partial: ", partial_b)
    print("Score C: ", score_c, caption)


# ======================================================================
#
#                           Brad T. Aagaard
#                        U.S. Geological Survey
#
# ======================================================================
#

import os
import logging
from importlib import import_module
import numpy

from . import analysis_utils
from . import gdalraster

class CostSavings(object):
    """Cost savings weighted by area and population.
    """

    def __init__(self, config):
        self.config = config
        return

    def compute(self, event, shakemap, alerts, shakingTime, populationDensity, magAlertThreshold, mmiAlertThreshold, plotAlertMaps=False):

        functionPath = self.config.get("mmi_predicted", "function").split(".")
        fn = getattr(import_module(".".join(functionPath[:-1])), functionPath[-1])

        shape = shakemap.data["mmi"].shape
        warningTimeZero = numpy.zeros((1,), dtype="timedelta64[us]")
        warningTime = gdalraster.NO_DATA_VALUE * 1.0e+6 * numpy.ones(shape, dtype="timedelta64[us]")
        mmiPred = gdalraster.NO_DATA_VALUE * numpy.ones(shape, numpy.float32)

        gmpe = self.config.get("mmi_predicted", "gmpe")
        gmice = self.config.get("mmi_predicted", "gmice")
        alertLatency = numpy.timedelta64(int(self.config.getfloat("alerts", "alert_latency_sec")*1.0e+3), "ms")
        if gmice == "default":
            gmice = shakemap.gmiceGrid

        thresholdReached = False
        for alert in alerts:
            alertTime = numpy.datetime64(alert["timestamp"]) + alertLatency

            if alertTime > numpy.max(shakingTime):
                # Skip alerts with no positive warning times in
                # domain. Changes in estimated earthquake location
                # could result in later alerts having positive warning
                # times.
                logging.getLogger(__name__).debug("Skipping alert version {ver} with no positive warning times.".format(ver=alert["version"]))
                continue
            if not thresholdReached and alert["magnitude"] < magAlertThreshold:
                continue
            else:
                if not thresholdReached:
                    wtime = analysis_utils.timedelta_to_seconds(alertTime - numpy.datetime64(event["origin_time"]))
                    msg = "Alert threshold reached at {tstamp}, {wtime:.1f}s after origin time.".format(tstamp=alertTime, wtime=wtime)
                    logging.getLogger(__name__).info(msg)
                    thresholdReached = True

            mmiPredCur = fn(alert, shakemap.data, gmpe, gmice)
            warningTimeCur = shakingTime - alertTime

            if plotAlertMaps:
                plotsDir = self.config.get("files", "plots_dir")
                if not os.path.isdir(plotsDir):
                    os.makedirs(plotsDir)
                filename = analysis_utils.analysis_event_label(self.config, self.eqId, magAlertThreshold, mmiAlertThreshold)+"_alert_snapshot.tiff"
                values = [
                    ("mmi_pred", mmiPredCur,),
                    ("warning_time", analysis_utils.timedelta_to_seconds(warningTimeCur),),
                ]
                gdalraster.write(filename, values, shakemap.num_lon(), shakemap.num_lat(), shakemap.spatial_ref(), shakemap.geo_transform())
                mapPanels = maps.MapPanels(self.config)
                mapPanels.load_data(event["event_id"], alert=alert)
                tafterOT = analysis_utils.timedelta_to_seconds(alertTime-numpy.datetime64(event["origin_time"]))
                mapPanels.mmi_warning_time(tafterOT)

            # Update alert time if greater than previous
            maskAlert = numpy.bitwise_and(warningTimeCur > warningTime, mmiPredCur >= mmiAlertThreshold)
            warningTime[maskAlert] = warningTimeCur[maskAlert]

            # Update predicted MMI if greater than previous AND
            # positive warning time. Assumes action will be taken if
            # alert threshold is reached (cannot be undone if later
            # updates reduce predicted MMI).
            maskMMI = numpy.bitwise_and(mmiPredCur > mmiPred, warningTimeCur >= warningTimeZero)
            mmiPred[maskMMI] = mmiPredCur[maskMMI]

        filename = "analysis_" + analysis_utils.analysis_event_label(self.config, event["event_id"], magAlertThreshold, mmiAlertThreshold) + ".tiff"
        metrics = self._cost(mmiPred, shakemap, warningTime, populationDensity, mmiAlertThreshold, filename)
        return metrics

    def _cost(self, mmiPred, shakemap, warningTime, populationDensity, mmiAlertThreshold, filename):
        """Compute cost savings metrics.
        """
        mmiObs = shakemap.data["mmi"]

        # Compute costNoEEW, costEEW, costPerfectEEW, costSavings
        objectPath = self.config.get("fragility_curves", "object").split(".")
        fragilityOptions = dict(self.config.items("fragility_curves"))
        fragilityOptions.pop("object")
        fragilityOptions.pop("label")
        fragilityOptions = {k: float(v) for k,v in fragilityOptions.items()}
        fragility = getattr(import_module(".".join(objectPath[:-1])), objectPath[-1])(**fragilityOptions)

        costDamage = fragility.cost_damage(mmiObs)
        costActionObs = fragility.cost_action(mmiObs)
        costNoEEW = costDamage
        costPerfectEEW = costDamage*(costDamage < costActionObs) + costActionObs*(costDamage >= costActionObs)
        costEEW = fragility.cost_action(mmiPred)*(mmiPred >= mmiAlertThreshold) + costDamage*(mmiPred < mmiAlertThreshold)

        pixelArea = shakemap.pixel_area(self.config.get("shakemap", "projection"))
        areaCostNoEEW = numpy.sum(pixelArea * costNoEEW)
        areaCostPerfectEEW = numpy.sum(pixelArea * costPerfectEEW)
        areaCostEEW = numpy.sum(pixelArea * costEEW)
        areaDamage = numpy.sum(pixelArea * (costDamage > 0.0))
        areaAlert = numpy.sum(pixelArea * (mmiPred >= mmiAlertThreshold))
        areaAlertPerfect = numpy.sum(pixelArea * (costDamage > costActionObs))

        popCostNoEEW = numpy.sum(populationDensity * pixelArea * costNoEEW)
        popCostPerfectEEW = numpy.sum(populationDensity * pixelArea * costPerfectEEW)
        popCostEEW = numpy.sum(populationDensity * pixelArea * costEEW)
        popDamage = numpy.sum(pixelArea * populationDensity * (costDamage > 0.0))
        popAlert = numpy.sum(pixelArea * populationDensity * (mmiPred >= mmiAlertThreshold))
        popAlertPerfect = numpy.sum(pixelArea * populationDensity * (costDamage > costActionObs))

        # Alert categories TN(0),FN(1),FP(2),TP(3)
        alertCategory = numpy.zeros(costDamage.shape)
        maskTN = numpy.bitwise_and(mmiPred < mmiAlertThreshold, costDamage < costActionObs)
        maskFN = numpy.bitwise_and(mmiPred < mmiAlertThreshold, costDamage >= costActionObs)
        maskFP = numpy.bitwise_and(mmiPred >= mmiAlertThreshold, costDamage < costActionObs)
        maskTP = numpy.bitwise_and(mmiPred >= mmiAlertThreshold, costDamage >= costActionObs)
        alertCategory = maskTN*0.0 + maskFN*1.0 + maskFP*2.0 + maskTP*3.0

        values = [
            ("mmi_obs", mmiObs,),
            ("mmi_pred", mmiPred,),
            ("warning_time", analysis_utils.timedelta_to_seconds(warningTime),),
            ("population_density", populationDensity,),
            ("cost_no_eew", costNoEEW,),
            ("cost_perfect_eew", costPerfectEEW,),
            ("cost_eew", costEEW,),
            ("alert_category", alertCategory,),
            ("pixel_area", pixelArea,),
            ]
        cacheDir = self.config.get("files", "analysis_cache_dir")
        if not os.path.isdir(cacheDir):
            os.makedirs(cacheDir)
        gdalraster.write(os.path.join(cacheDir, filename), values, shakemap.num_lon(), shakemap.num_lat(), shakemap.spatial_ref(), shakemap.geo_transform())

        metrics = {
            "area_damage": areaDamage,
            "area_alert": areaAlert,
            "area_alert_perfect": areaAlertPerfect,
            "area_costsavings_eew": areaCostNoEEW - areaCostEEW,
            "area_costsavings_perfecteew": areaCostNoEEW - areaCostPerfectEEW,
            "population_damage": popDamage,
            "population_alert": popAlert,
            "population_alert_perfect": popAlertPerfect,
            "population_costsavings_eew": popCostNoEEW - popCostEEW,
            "population_costsavings_perfecteew": popCostNoEEW - popCostPerfectEEW,
            }
        return metrics

# End of file