#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
SRI International社の Inter Gamut パターンを改良して、
Gamut の境界を検出するパターンを作成する。
"""

# 外部ライブラリのインポート
import os
import numpy as np
import cv2

# 自作ライブラリのインポート
import test_pattern_generator2 as tpg
import color_space as cs
import transfer_functions as tf
from DrawGamutPattern import DrawGamutPattern
from DrawChromaticityDiagram import DrawChromaticityDiagram
from DrawInformation import DrawInformation
from CalcParameters import CalcParameters

mini_primaries = [[0.45, 0.25], [0.30, 0.45], [0.25, 0.20], [0.45, 0.25]]

BASE_PARAM = {
    'revision': 1,
    'inner_sample_num': 4,
    'outer_sample_num': 4,
    'hue_devide_num': 6,
    'img_width': 1920,
    'img_height': 1080,
    'pattern_space_rate': 0.71,
    'inner_gamut_name': 'ITU-R BT.709',
    'outer_gamut_name': 'ITU-R BT.2020',
    'inner_primaries': np.array(tpg.get_primaries(cs.BT709)[0]),
    'outer_primaries': np.array(tpg.get_primaries(cs.BT2020)[0]),
    'transfer_function': tf.SRGB,
    'background_luminance': 2,
    'reference_white': 100
}

class GamutEdgeChecker:
    """
    ## 概要

    Gamutのエッジをチェックする。

    ## 全体の処理の概要

    ```
    gamut_edge_checker = GamutEdgeChecker(base_param)
    gamut_edge_checker.make()
    gamut_edge_checker.preview()
    gamut_edge_checker.save()
    ```

    ## Gec.make() の概要

    ```
    def make(self):
        self.make_base_layer()  // 大元の背景画像を準備

        calc_param = CalcParameters(self.base_param)
        draw_param = calc_param.calc_parameters()

        draw_pattern = DrawGamutPattern(draw_param, self.img)
        draw_pattern.draw_gamut_tile_pattern()

        draw_diagram = DrawChromaticityDiagram(draw_param, self.img)
        draw_diagram.draw_chromaticity_diagram()

        text_info = self.make_text_information()
        draw_information = DrawInformation(text_imfo, self.img)
        draw_information.draw_information()

        self.apply_oetf()
    ```

    ## self.make_base_param() の吐き出す値

    ```
    typedef struct{
        int revision;
        int inner_sample_num;  // 内側の描画点の数
        int outer_sample_num;  // 外側の描画点の数
        int hue_devide_num;  // 色相方向の分割数。原則4固定。
        int img_width;
        int img_height;
        char *inner_gamut_name;  // 内側の Gamut名
        char *outer_gamut_name;  // 外側の Gamut名
        double inner_primaries[3][3];  // 内側のxy色度座標
        double outer_primaries[3][3];  // 外側のxy色度座標
        char *transfer_function;  // OETF の指定
        int reference_white;  // ref white の設定。単位は [cd/m2]。
    }base_param;
    ```

    ## self.make_text_information() の吐き出す値

    ```
    typedef struct{
        int diagram_width;
        int diagram_height;
    }text_info;
    ```

    ## calc_param.calc_parameters() の吐き出す draw_param 値

    typedef struct{
        double inner_xyY[12][inner_sample_num][3];
        double outer_xyY[12][outer_sample_num][3];
        double innnr_ref_xyY[12][inner_sample_num][3];
        double outer_ref_xyY[12][inner_sample_num][3];
        double min_large_y[12];  // inner_xy, outer_xy の largeY最小値。
                                 // これに合わせて xyY to RGB 変換を行う
    }draw_param // 12 は 3(RGB) * 4(hue_devide_num) から算出
    """
    def __init__(self, base_param=BASE_PARAM):
        self.base_param = base_param

    def make(self):
        """
        画像生成
        """
        self.make_base_layer()
        calc_param = CalcParameters(self.base_param)
        draw_param = calc_param.calc_parameters()
        draw_pattern = DrawGamutPattern(self.base_param, draw_param, self.img)
        draw_pattern.draw_gamut_tile_pattern()
        draw_diagram = DrawChromaticityDiagram(
            self.base_param, draw_param, self.img)
        draw_diagram.draw_chromaticity_diagram()
        diagram_width, diagram_height =\
            draw_diagram.get_diagram_widgh_height()
        draw_information = DrawInformation(
            self.base_param, draw_param, self.img,
            diagram_width, diagram_height)
        draw_information.draw_information()
        self.apply_oetf()

    def int(self, x):
        return int(x + 0.5)

    def make_base_layer(self):
        """
        ベースとなる背景画像を生成。
        """
        # 大枠準備
        width = self.base_param['img_width']
        height = self.base_param['img_height']
        self.img = np.zeros((height, width, 3))

        # Gamut パターン配置場所のBG Colorを設定
        pattern_space_rate = self.base_param['pattern_space_rate']
        background_luminance = self.base_param['background_luminance']
        gamut_area_width = int(width * pattern_space_rate)
        bg_img = np.ones((height, gamut_area_width, 3)) * background_luminance
        self.img[:, :gamut_area_width, :] = bg_img

    def apply_oetf(self):
        oetf_name = self.base_param['transfer_function']
        self.img = tf.oetf_from_luminance(self.img, oetf_name)
        if np.sum(self.img > 1.0) > 0:
            print("warning. over flow")
        elif np.sum(self.img < 0.0) > 0:
            print("warning. under flow")
        self.img = np.clip(self.img, 0, 1)
        self.img = np.uint16(np.round(self.img * 0xFFFF))

    def preview(self):
        tpg.preview_image(self.img)

    def save(self):
        cv2.imwrite("test.tiff", self.img[:, :, ::-1])

    def print_output_xyY_header(self):
        h_num = self.base_param['hue_devide_num'] * 3
        strings = "No" + ",x,y,Y" * h_num
        print(strings)

    def output_xyY(self):
        """
        デバッグとか他のプロジェクトでの流用とかを見越して
        xyY 情報を出力する。
        """
        calc_param = CalcParameters(self.base_param)
        draw_param = calc_param.calc_parameters()
        self.print_output_xyY_header()
        s = ",{:.6f},{:.6f},{:.6f}"
        counter = 0
        for idx in range(self.base_param['inner_sample_num']):
            print(counter, end="")
            for h_idx in range(self.base_param['hue_devide_num'] * 3):
                print(s.format(draw_param['inner_xyY'][h_idx][idx][0],
                               draw_param['inner_xyY'][h_idx][idx][1],
                               draw_param['inner_xyY'][h_idx][idx][2]),
                      end="")
            print("")
            counter += 1
        for idx in range(self.base_param['outer_sample_num']):
            print(counter, end="")
            for h_idx in range(self.base_param['hue_devide_num'] * 3):
                print(s.format(draw_param['outer_xyY'][h_idx][idx][0],
                               draw_param['outer_xyY'][h_idx][idx][1],
                               draw_param['outer_xyY'][h_idx][idx][2]),
                      end="")
            print("")
            counter += 1

def main_func():
    gamut_edge_checker = GamutEdgeChecker(base_param=BASE_PARAM)
    # gamut_edge_checker.make()
    # gamut_edge_checker.preview()
    # gamut_edge_checker.save()
    gamut_edge_checker.output_xyY()

if __name__ == '__main__':
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    main_func()


from dataclasses import dataclass, field
from typing import Sequence, List, Optional, cast, Any, TYPE_CHECKING
import lightweaver.constants as Const
import numpy as np
from .atomic_table import PeriodicTable
from .barklem import Barklem

if TYPE_CHECKING:
    from .atomic_model import AtomicLine
    from .atmosphere import Atmosphere
    from .atomic_set import SpeciesStateTable

@dataclass
class LineBroadeningResult:
    '''
    Result expected from instances of `LineBroadening.broaden`.
    '''
    natural: np.ndarray
    Qelast: np.ndarray
    other: Optional[List] = None

@dataclass
class LineBroadener:
    '''
    Base class for broadening implementations. To be used if your broadener
    does something special and can't just return an array.
    '''
    def __repr__(self):
        raise NotImplementedError

    def setup(self, line: 'AtomicLine'):
        pass

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> Any:
        raise NotImplementedError

@dataclass
class StandardLineBroadener(LineBroadener):
    '''
    Standard base class for broadening implementations. Unless you need to do
    something weird, inherit from this one.
    '''
    def __repr__(self):
        raise NotImplementedError

    def setup(self, line: 'AtomicLine'):
        pass

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        raise NotImplementedError

@dataclass
class LineBroadening:
    '''
    Standard component of AtomicLine to compute the broadening parameters in
    a flexible way.

    For most Voigt-like situations, this should be usable without
    modifications, but this class can be inherited from to make substantial
    modifications.

    Parameters
    ----------
    natural : list of StandardLineBroadener
        List of broadening terms that are not elastic collisions (separated for PRD).
    elastic : list of StandardLineBroadener
        List of elastic broadening terms.
    other : list of LineBroadener, optional
        List of other broadening terms, not used by the VoigtLine by default,
        but existing to provide _options_ (default: None)
    '''
    natural: List[StandardLineBroadener]
    elastic: List[StandardLineBroadener]
    other: Optional[List[LineBroadener]] = None

    def __repr__(self):
        otherStr = '' if self.other is None else ', other=%s' % repr(self.other)
        s = 'LineBroadening(natural=%s, elastic=%s%s)' % (repr(self.natural), repr(self.elastic), otherStr)
        return s

    def __post_init__(self):
        if len(self.natural) == 0 and len(self.elastic) == 0:
            raise ValueError('No standard broadening terms provided to LineBroadening')

    def setup(self, line: 'AtomicLine'):
        b: LineBroadener
        for b in self.natural:
            b.setup(line)

        for b in self.elastic:
            b.setup(line)

        if self.other is not None:
            for b in self.other:
                b.setup(line)

    @staticmethod
    def sum_broadening_list(broadeners: List[StandardLineBroadener], atmos: 'Atmosphere',
                            eqPops: 'SpeciesStateTable') -> Optional[np.ndarray]:
        '''
        Sums a list of StandardLineBroadeners.
        '''
        if len(broadeners) == 0:
            return None

        result = broadeners[0].broaden(atmos, eqPops)
        for b in broadeners[1:]:
            result += b.broaden(atmos, eqPops)
        return result

    @staticmethod
    def compute_other_broadening(broadeners: Optional[List[LineBroadener]],
                                 atmos: 'Atmosphere',
                                 eqPops: 'SpeciesStateTable') -> Optional[List]:
        '''
        Returns a list of the computed broadening terms.
        '''

        if broadeners is None:
            return None
        if len(broadeners) == 0:
            return None

        result = [b.broaden(atmos, eqPops) for b in broadeners]
        return result

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> LineBroadeningResult:
        '''
        Computes the broadening, this function is called by the AtomicLine object.
        '''
        natural = self.sum_broadening_list(self.natural, atmos, eqPops)
        Qelast = self.sum_broadening_list(self.elastic, atmos, eqPops)

        others = self.compute_other_broadening(self.other, atmos, eqPops)

        if natural is None:
            natural = np.zeros_like(Qelast)
        elif Qelast is None:
            Qelast = np.zeros_like(natural)

        return LineBroadeningResult(natural=natural,
                                    Qelast=Qelast, other=others)

@dataclass(eq=False)
class VdwApprox(StandardLineBroadener):
    '''
    Base class for van der Waals approximation using a list of coefficients
    (vals).
    '''
    vals: Sequence[float]
    line: 'AtomicLine' = field(init=False)

    def setup(self, line: 'AtomicLine'):
        self.line = line

    def __repr__(self):
        s = '%s(vals=%s)' % (type(self).__name__, repr(self.vals))
        return s

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        if self.vals != other.vals:
            return False

        try:
            if self.line != other.line:
                return False
        except:
            pass

        return True

@dataclass(eq=False, repr=False)
class VdwUnsold(VdwApprox):
    '''
    Implementation of the Unsold method for van der Waals broadening.
    '''
    def setup(self, line: 'AtomicLine'):
        self.line = line
        if len(self.vals) != 2:
            raise ValueError('VdwUnsold expects 2 coefficients (%s)' % repr(line))

        Z = line.jLevel.stage + 1
        cont = line.overlyingContinuumLevel

        deltaR = (Const.ERydberg / (cont.E_SI - line.jLevel.E_SI))**2 \
                 - (Const.ERydberg / (cont.E_SI - line.iLevel.E_SI))**2
        fourPiEps0 = 4.0 * np.pi * Const.Epsilon0
        self.C625 = (2.5 * Const.QElectron**2 / fourPiEps0 * Const.ABarH / fourPiEps0 \
                     * 2 * np.pi * (Z * Const.RBohr)**2 / Const.HPlanck * deltaR)**0.4

        element = line.atom.element

        self.vRel35He = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
                         * (1.0 + element.mass / PeriodicTable[2].mass))**0.3
        self.vRel35H = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
                         * (1.0 + element.mass / PeriodicTable[1].mass))**0.3

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        heAbund = eqPops.abundance[PeriodicTable[2]]
        cross = 8.08 * (self.vals[0] * self.vRel35H \
                             + self.vals[1] * heAbund * self.vRel35He) * self.C625
        nHGround = eqPops['H'][0, :]
        broad = cross * atmos.temperature**0.3 * nHGround
        return broad

@dataclass(eq=False, repr=False)
class VdwBarklem(VdwApprox):
    '''
    Implementation of the Barklem method for van der Waals broadening.
    '''
    def setup(self, line: 'AtomicLine'):
        self.line = line
        if len(self.vals) != 2:
            raise ValueError('VdwBarklem expects 2 coefficients (%s)' % (repr(line)))
        newVals = Barklem.get_active_cross_section(line.atom, line, self.vals)

        self.barklemVals = newVals

        Z = line.jLevel.stage + 1
        cont = line.overlyingContinuumLevel

        deltaR = (Const.ERydberg / (cont.E_SI - line.jLevel.E_SI))**2 \
                 - (Const.ERydberg / (cont.E_SI - line.iLevel.E_SI))**2
        fourPiEps0 = 4.0 * np.pi * Const.Epsilon0
        self.C625 = (2.5 * Const.QElectron**2 / fourPiEps0 * Const.ABarH / fourPiEps0 \
                     * 2 * np.pi * (Z * Const.RBohr)**2 / Const.HPlanck * deltaR)**0.4

        element = line.atom.element

        self.vRel35He = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
                         * (1.0 + element.mass / PeriodicTable[2].mass))**0.3

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        heAbund = eqPops.abundance[PeriodicTable[2]]
        nHGround = eqPops['H'][0, :]
        cross = 8.08 * self.barklemVals[2] * heAbund * self.vRel35He * self.C625

        broad = self.barklemVals[0] * atmos.temperature**(0.5*(1.0-self.barklemVals[1])) \
                 + cross * atmos.temperature**0.3
        broad *= nHGround
        return broad

@dataclass(eq=False)
class RadiativeBroadening(StandardLineBroadener):
    '''
    Simple constant radiative broadening with coefficient gamma.
    '''
    gamma: float
    line: 'AtomicLine' = field(init=False)

    def setup(self, line: 'AtomicLine'):
        self.line = line

    def __repr__(self):
        s = '%s(gamma=%g)' % (type(self).__name__, self.gamma)
        return s

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        if self.gamma != other.gamma:
            return False

        try:
            if self.line != other.line:
                return False
        except:
            pass

        return True

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        return np.ones_like(atmos.temperature) * self.gamma

@dataclass
class QuadraticStarkBroadening(StandardLineBroadener):
    '''
    Lindholm theory result for Quadratic Stark broadening by electrons and
    singly ionised particles.
    Follows HM2014 pp. 238-239, uses C4 from Traving 1960 via RH.
    '''
    coeff: float
    line: 'AtomicLine' = field(init=False)

    def __repr__(self):
        s = '%s(coeff=%g)' % (type(self).__name__, self.coeff)
        return s

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        if self.coeff != other.coeff:
            return False

        try:
            if self.line != other.line:
                return False
        except:
            pass

        return True

    def setup(self, line: 'AtomicLine'):
        self.line = line
        weight = line.atom.element.mass
        C = 8.0 * Const.KBoltzmann / (np.pi * Const.Amu * weight)
        Cm = (1.0 + weight / (Const.MElectron / Const.Amu))**(1.0/6.0)
        # NOTE(cmo): 28.0 is average atomic weight
        Cm += (1.0 + weight / (28.0))**(1.0/6.0)
        self.C = C
        self.Cm = Cm

        Z = line.iLevel.stage + 1
        cont = line.overlyingContinuumLevel

        E_Ryd = Const.ERydberg / (1.0 + Const.MElectron / (weight * Const.Amu))
        neff_l = Z * np.sqrt(E_Ryd / (cont.E_SI - line.iLevel.E_SI))
        neff_u = Z * np.sqrt(E_Ryd / (cont.E_SI - line.jLevel.E_SI))

        C4 = Const.QElectron**2 / (4.0 * np.pi * Const.Epsilon0) \
            * Const.RBohr \
            * (2.0 * np.pi * Const.RBohr**2 / Const.HPlanck) / (18.0 * Z**4) \
            * ((neff_u * (5.0 * neff_u**2 + 1.0))**2 \
                - (neff_l * (5.0 * neff_l**2 + 1.0))**2)
        self.cStark23 = 11.37 * (self.coeff * C4)**(2.0/3.0)

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        vRel = (self.C * atmos.temperature)**(1.0/6.0) * self.Cm
        stark = self.cStark23 * vRel * atmos.ne
        return stark

@dataclass
class MultiplicativeStarkBroadening(StandardLineBroadener):
    '''
    Simple expression for multiplicative Stark broadening, assumes that this
    can be expresed as a constant * ne.
    '''
    coeff: float

    def __repr__(self):
        s = '%s(coeff=%g)' % (type(self).__name__, self.coeff)
        return s

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        if self.coeff != other.coeff:
            return False

        return True

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        return self.coeff * atmos.ne # type: ignore

@dataclass
class HydrogenLinearStarkBroadening(StandardLineBroadener):
    '''
    Linear Stark broadening for the case of Hydrogen from Sutton 1978 (like
    RH).
    '''
    line: 'AtomicLine' = field(init=False)

    def __repr__(self):
        s = '%s()' % type(self).__name__
        return s

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        try:
            if self.line != other.line:
                return False
        except:
            pass

        return True

    def setup(self, line: 'AtomicLine'):
        self.line = line

        if line.atom.element.Z != 1:
            raise ValueError('HydrogenicLinearStarkBroadening applied to non-Hydrogen line')

    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
        '''
        The function that is called by LineBroadening.

        Parameters
        ----------
        atmos : Atmosphere
            The atmosphere in which to compute the broadening.
        eqPops : SpeciesStateTable
            The populations to use for computing the broadening.

        Returns
        -------
        broad : np.ndarray
            An array detailing the broadening at each location in the
            atmosphere [Nspace].
        '''
        nUpper = int(np.round(np.sqrt(0.5*self.line.jLevel.g)))
        nLower = int(np.round(np.sqrt(0.5*self.line.iLevel.g)))

        a1 = 0.642 if nUpper - nLower == 1 else 1.0
        C = a1 * 0.6 * (nUpper**2 - nLower**2) * Const.CM_TO_M**2
        GStark = C * atmos.ne**(2.0/3.0)
        return GStark


# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from collections import OrderedDict

from ....serialization.serializables import (
    Int64Field,
    BoolField,
    Int32Field,
    Float64Field,
)
from ....utils import pd_release_version
from ...utils import validate_axis
from ..core import Window

_default_min_period_1 = pd_release_version >= (1, 1, 0)
_pd_1_3_repr = pd_release_version >= (1, 3, 0)

class EWM(Window):
    _alpha = Float64Field("alpha")
    _min_periods = Int64Field("min_periods")
    _adjust = BoolField("adjust")
    _ignore_na = BoolField("ignore_na")
    _axis = Int32Field("axis")

    def __init__(
        self, alpha=None, min_periods=None, adjust=None, ignore_na=None, axis=None, **kw
    ):
        super().__init__(
            _alpha=alpha,
            _min_periods=min_periods,
            _adjust=adjust,
            _ignore_na=ignore_na,
            _axis=axis,
            **kw
        )

    @property
    def alpha(self):
        return self._alpha

    @property
    def min_periods(self):
        return self._min_periods

    @property
    def adjust(self):
        return self._adjust

    @property
    def ignore_na(self):
        return self._ignore_na

    @property
    def axis(self):
        return self._axis

    @property
    def params(self):
        p = OrderedDict()
        for k in ["alpha", "min_periods", "adjust", "ignore_na", "axis"]:
            p[k] = getattr(self, k)
        return p

    def __call__(self, df):
        return df.ewm(**self.params)

    def _repr(self, params):
        com = 1.0 / params.pop("alpha") - 1
        params["com"] = int(com) if _pd_1_3_repr and com == math.floor(com) else com
        try:
            params.move_to_end("com", last=False)
        except AttributeError:  # pragma: no cover
            pass
        return super()._repr(params)

    def _repr_name(self):
        try:
            from pandas.core.window import ExponentialMovingWindow  # noqa: F401

            return "ExponentialMovingWindow"
        except ImportError:  # pragma: no cover
            return "EWM"

    def aggregate(self, func):
        from .aggregation import DataFrameEwmAgg

        params = self.params
        params["alpha_ignore_na"] = params.pop("ignore_na", False)
        params["validate_columns"] = False
        op = DataFrameEwmAgg(func=func, **params)
        return op(self)

    agg = aggregate

    def mean(self):
        return self.aggregate("mean")

    def var(self):
        return self.aggregate("var")

    def std(self):
        return self.aggregate("std")

def ewm(
    obj,
    com=None,
    span=None,
    halflife=None,
    alpha=None,
    min_periods=0,
    adjust=True,
    ignore_na=False,
    axis=0,
):
    r"""
    Provide exponential weighted functions.

    Parameters
    ----------
    com : float, optional
        Specify decay in terms of center of mass,
        :math:`\alpha = 1 / (1 + com),\text{ for } com \geq 0`.
    span : float, optional
        Specify decay in terms of span,
        :math:`\alpha = 2 / (span + 1),\text{ for } span \geq 1`.
    halflife : float, optional
        Specify decay in terms of half-life,
        :math:`\alpha = 1 - exp(log(0.5) / halflife),\text{for} halflife > 0`.
    alpha : float, optional
        Specify smoothing factor :math:`\alpha` directly,
        :math:`0 < \alpha \leq 1`.
    min_periods : int, default 0
        Minimum number of observations in window required to have a value
        (otherwise result is NA).
    adjust : bool, default True
        Divide by decaying adjustment factor in beginning periods to account
        for imbalance in relative weightings
        (viewing EWMA as a moving average).
    ignore_na : bool, default False
        Ignore missing values when calculating weights;
        specify True to reproduce pre-0.15.0 behavior.
    axis : {0 or 'index', 1 or 'columns'}, default 0
        The axis to use. The value 0 identifies the rows, and 1
        identifies the columns.

    Returns
    -------
    DataFrame
        A Window sub-classed for the particular operation.

    See Also
    --------
    rolling : Provides rolling window calculations.
    expanding : Provides expanding transformations.

    Notes
    -----
    Exactly one of center of mass, span, half-life, and alpha must be provided.

    Allowed values and relationship between the parameters are specified in the
    parameter descriptions above; see the link at the end of this section for
    a detailed explanation.

    When adjust is True (default), weighted averages are calculated using
    weights (1-alpha)**(n-1), (1-alpha)**(n-2), ..., 1-alpha, 1.

    When adjust is False, weighted averages are calculated recursively as:

       weighted_average[0] = arg[0];
       weighted_average[i] = (1-alpha)*weighted_average[i-1] + alpha*arg[i].

    When ignore_na is False (default), weights are based on absolute positions.
    For example, the weights of x and y used in calculating the final weighted
    average of [x, None, y] are (1-alpha)**2 and 1 (if adjust is True), and
    (1-alpha)**2 and alpha (if adjust is False).

    When ignore_na is True (reproducing pre-0.15.0 behavior), weights are based
    on relative positions. For example, the weights of x and y used in
    calculating the final weighted average of [x, None, y] are 1-alpha and 1
    (if adjust is True), and 1-alpha and alpha (if adjust is False).

    More details can be found at
    https://pandas.pydata.org/pandas-docs/stable/user_guide/computation.html#exponentially-weighted-windows

    Examples
    --------
    >>> import numpy as np
    >>> import mars.dataframe as md
    >>> df = md.DataFrame({'B': [0, 1, 2, np.nan, 4]})
    >>> df.execute()
         B
    0  0.0
    1  1.0
    2  2.0
    3  NaN
    4  4.0
    >>> df.ewm(com=0.5).mean().execute()
              B
    0  0.000000
    1  0.750000
    2  1.615385
    3  1.615385
    4  3.670213
    """
    axis = validate_axis(axis, obj)

    decay_count = 0
    for arg in (com, span, halflife, alpha):
        if arg is not None:
            decay_count += 1

    if decay_count == 0:
        raise ValueError("Must pass one of comass, span, halflife, or alpha")
    if decay_count > 1:
        raise ValueError("comass, span, halflife, and alpha are mutually exclusive")

    if com is not None:
        if com < 0:
            raise ValueError("comass must satisfy: comass >= 0")
        alpha = 1.0 / (1 + com)
    elif span is not None:
        if span < 1:
            raise ValueError("span must satisfy: span >= 1")
        alpha = 2.0 / (1 + span)
    elif halflife is not None:
        if halflife <= 0:
            raise ValueError("halflife must satisfy: halflife > 0")
        alpha = 1.0 - math.exp(math.log(0.5) / halflife)
    if alpha <= 0 or alpha > 1:
        raise ValueError("alpha must satisfy: 0 < alpha <= 1")

    if not adjust and not ignore_na:
        raise NotImplementedError(
            "adjust == False when ignore_na == False not implemented"
        )
    if axis == 1:
        raise NotImplementedError("axis other than 0 is not supported")

    if alpha == 1:
        return obj.expanding(min_periods=min_periods, axis=axis)

    if _default_min_period_1:
        min_periods = min_periods or 1

    return EWM(
        input=obj,
        alpha=alpha,
        min_periods=min_periods,
        adjust=adjust,
        ignore_na=ignore_na,
        axis=axis,
    )


import os
import sys
import numpy as np
import math
from matplotlib import pyplot as plt
import six
from pyscirsc.gromacs.molecule import correct_pbc, correct_box
from pyscirsc.counting.msd.msd_common import MSDException
from pyscirsc.gromacs.band import Band
from pmx.model import Model
from pmx.xtc import Trajectory
from pmx.ndx import IndexFile
from pyscirsc.gromacs.index import indexes
from pyscirsc.gromacs.trajectory import iterate
from pyscirsc.utils import mkdir_p

class MSDSimplifiedDeviation:

    def __init__(self,
                 atom_group, bands, restarts,
                 computation_ttl=[10,100],
                 verbose=False):
        """
        :param atom_group: some AtomSelection as in pmx.atomselection
        :param band:       pyscirsc.gromacs.Band object
        :param restarts:   list of frame numbers to start new msd computation
                           (use msd select_starts ... for this purpose)
        :param verbose:    [default:False] print some info to stdout in process
        """
        # input
        # for i, b in enumerate(bands):
        #     print i, b.left, b.right
        # for i, r in enumerate(restarts):
        #     print i, r
        self.bands = bands
        self.group = atom_group
        self.restarts = restarts
        self.computation_ttl = computation_ttl
        self.starts = [map(lambda x: x + computation_ttl[0], r) for r in self.restarts]
        # self.finishes = [map(lambda x: x + computation_ttl[1], r) for r in self.restarts]
        self.finishes = [[] for r in self.restarts]
        self.current_restart = [0 for band in self.bands]
        self.current_start = [0 for band in self.bands]
        self.current_finish = [0 for band in self.bands]
        self.xs = [[] for band in self.bands]
        self.ys = [[] for band in self.bands]
        self.zs = [[] for band in self.bands]
        self.xf = [[] for band in self.bands]
        self.yf = [[] for band in self.bands]
        self.zf = [[] for band in self.bands]
        self.ref_positions = [[] for band in self.bands]
        self.times = []
        self.frames = []
        self.traj_len = 0
        self.verbose = True
        self.prev_pos = [(a.x[0], a.x[1], a.x[2]) for a in self.group]
        if not restarts:
            raise MSDException("No restart points given!")

    def hello(self):
        s = ""
        s += "# hello, its MSDSimplifiedDeviation!\n"
        s += "#    band:          %s\n" % str(self.band)
        s += "#    atom_group:    %s\n" % str([a.x for a in self.group])
        s += "#    restarts:      %s\n" % self.restarts
        s += "#    reference_ttl: %s\n" % self.reference_ttl
        return s

    def count(self, data, frame, fc):
        """
        Post-trajectory-update hook to be send to
        pyscirsc.gromacs.trajectory.iterate(...)
        """
        self.times.append(frame.time)
        self.frames.append(fc)
        self.traj_len += 1

        if (True or self.verbose) and fc % 100 == 0:
            print "# at frame: ", fc

        # if fc not in [self.restarts[r] for r in self.current_restart] and \
        #     fc not in [self.finishes[r] for r in self.current_finish]:
        #     return

        inband = []
        for band in self.bands:

            # set inband[i] = true if atom position is inside band borders
            # we test BEFORE removing pbc, as after correct_pbc it may go to
            # other periodic image... (we assume xtc trajectory is confined to BOX)
            inband.append(band.mask(self.group))

        # now correct periodicity (move atoms to their most probable position)
        box = correct_box(data)
        no_pbc = [correct_pbc(o, a.x, box) for o, a in zip(self.prev_pos, self.group)]

        for bandidx, band in enumerate(self.bands):
            # if fc != self.restarts[bandidx][self.current_restart[bandidx]] and \
            #     fc != [self.finishes[bandidx][self.current_finish[bandidx]]]:
            #     continue

            # append new restart position from list if
            # it was selected for this frame.
            if fc == self.restarts[bandidx][self.current_restart[bandidx]]:  # at defined intervals
                self.ref_positions[bandidx].append({
                    'positions': no_pbc,  # original positions of the atoms
                    'inband': inband[bandidx],     # mask for the atoms in band
                })
                # this guarantees is always a valid index in restarts
                if self.current_restart[bandidx] < len(self.restarts[bandidx]) - 1:
                    self.current_restart[bandidx] += 1

            if fc == self.starts[bandidx][self.current_start[bandidx]] or \
                (self.current_finish[bandidx] < len(self.finishes[bandidx]) and fc == self.finishes[bandidx][self.current_finish[bandidx]]):
                # at defined intervals

                if fc == self.starts[bandidx][self.current_start[bandidx]]:
                    ref = self.ref_positions[bandidx][self.current_start[bandidx]]
                else:
                    ref = self.ref_positions[bandidx][self.current_finish[bandidx]]

                x, y, z, N = 0.0, 0.0, 0.0, 0.0
                common_test = [u and v for u, v in zip(inband[bandidx], ref['inband'])]
                starts = ref['positions']
                for atom, orig, common in zip(no_pbc, starts, common_test):
                    if common:  # if atom is both in ref position and current inband
                        dx = (atom[0] - orig[0]) ** 2
                        dy = (atom[1] - orig[1]) ** 2
                        dz = (atom[2] - orig[2]) ** 2
                        x += dx
                        y += dy
                        z += dz
                        N += 1.0

                x /= 100.  # A^2 -> nm^2
                y /= 100.  # A^2 -> nm^2
                z /= 100.  # A^2 -> nm^2

                if N:
                    if fc == self.starts[bandidx][self.current_start[bandidx]]:
                        self.xs[bandidx].append(x / N)
                        self.ys[bandidx].append(y / N)
                        self.zs[bandidx].append(z / N)
                        print "appending start at", fc
                        self.finishes[bandidx].append(fc - self.computation_ttl[0] + self.computation_ttl[1])
                    else:
                        self.xf[bandidx].append(x / N)
                        self.yf[bandidx].append(y / N)
                        self.zf[bandidx].append(z / N)
                        print "finished at", fc

                else:
                    if any(inband[bandidx]):
                        if self.verbose:
                            msg = "# EMPTY INTERSECTION at frame={0} time={1} bandidx={2}"
                            print msg.format(fc, frame.time, bandidx)
                    else:
                        if self.verbose:
                            msg = "# EMPTY INBAND at frame={0} time={1} bandidx={2}"
                            print msg.format(fc, frame.time, bandidx)

                if fc == self.starts[bandidx][self.current_start[bandidx]]:
                    if self.current_start[bandidx] < len(self.starts[bandidx]) - 1:
                        self.current_start[bandidx] += 1
                else:
                    # this guarantees is always a valid index in finishes
                    if self.current_finish[bandidx] < len(self.finishes[bandidx]) - 1:
                        self.current_finish[bandidx] += 1

        # remember for next frame PBC correction
        self.prev_pos = no_pbc

################################################################################
# script to be used in msd.py                                                  #
################################################################################

# main function
def main(runfile_path=None):
    # You can adjust those variables
    pdb_filename   = "odnPOPC_001.pdb"  # pdb file with model
    xtc_filename   = "odnPOPC_test.xtc"  # trajectory file
    begin_frame    = 0     # start frame of the trajectory (in frames, not ps!)
    end_frame      = 1000  # end frame (if -1 then all) (in frames, not ps!)
    output_dir     = "data/output"  # output in current directory
    group_resname  = "OG"  # residue name of the selected group to compute msd
    reference_ttl  = 200   # length of msd calculation for each restart (not used at all in this script, but in msd_computation later)
    min_stay       = 5     # restart points only for atoms in band for that many consecutive frames
    min_gap        = 25    # minimal gap between restarts
    max_gap        = 50    # maximal gap between restarts
    band_dir       = 2     # 0 = x, 1 = y, 2 = z, z is default
    # band separators:
    separators = [  # computed externally (in nm!)
        0.0,         0.85879874,  1.51332498,  # as defined by Ela Plesner
        4.34488884,  5.00243812,  7.38666648,  # see Diff_tables.doc
        8.04406502, 10.8815546,  11.5530146,   # first must be 0, last is some
        25.0                                   # big number greater than the box
    ]
    separators = map(lambda x: 10 * x, separators)  # make it angstroems

    starts = {
        0: [15, 69, 97, 123, 164, 200, 229, 279, 329, 381, 408, 440, 540, 576,
            609, 636, 670, 695, 721, 749, 775, 803, 829, 869, 919, 964, 990],
        1: [1, 30, 60, 92, 120, 156, 187, 229, 256, 283, 315, 341, 381, 408,
            440, 467, 494, 520, 549, 575, 609, 634, 667, 692, 724, 749, 775,
            803, 831, 864, 894, 920, 953, 980],
        2: [1, 26, 51, 77, 104, 129, 154, 182, 209, 236, 273, 303, 330, 358,
            384, 409, 436, 461, 490, 515, 540, 576, 609, 635, 662, 694, 723,
            749, 775, 800, 829, 861, 886, 914, 939, 964, 990],
        3: [1, 51, 97, 123, 164, 196, 222, 247, 283, 329, 379, 408, 440, 490,
            540, 576, 609, 636, 670, 695, 745, 775, 803, 853, 886, 936, 964,
            990],
        4: [1, 51, 97, 123, 164, 200, 229, 279, 329, 379, 408, 440, 490, 540,
            576, 609, 636, 670, 695, 745, 775, 803, 853, 886, 936, 964, 990],
        5: [1, 38, 69, 97, 123, 164, 200, 229, 279, 306, 335, 360, 391, 419,
            446, 496, 535, 560, 595, 636, 662, 694, 724, 749, 775, 803, 837,
            863, 893, 921, 957, 983],
        6: [1, 26, 52, 81, 108, 140, 168, 195, 228, 254, 279, 309, 334, 365,
            391, 425, 450, 475, 500, 530, 557, 593, 627, 653, 679, 705, 730,
            755, 785, 811, 837, 864, 889, 914, 940, 965, 990],
        7: [1, 51, 97, 123, 164, 200, 229, 279, 329, 379, 408, 440, 490, 540,
            576, 609, 636, 670, 695, 720, 749, 775, 803, 853, 886, 936, 964,
            990],
        8: [1, 51, 89, 123, 165, 196, 225, 266, 316, 348, 377, 408, 435, 470,
            495, 524, 560, 609, 646, 679, 725, 764, 814, 864, 914, 964, 992],
    }

    print "###################################################"
    if runfile_path:
                # this is not nice, but powerfull..
        d = dict()
        # exec (open(runfile_path).read(), d, d)
        exec open(runfile_path).read() in d, d
        # TODO: grrrrrrrrrrr
        # TODO: because of the BUG I need to do this this way, grrrrrr....
        # TODO: make it a local config, (e.g. Main class with main() method)
        # TODO: and stroe config in self -> then can automatically assign in loop
        pdb_filename  = d.get('pdb_filename', pdb_filename)
        xtc_filename  = d.get('xtc_filename', xtc_filename)
        begin_frame   = d.get('begin_frame', begin_frame)
        end_frame     = d.get('end_frame', end_frame)
        output_dir    = d.get('output_dir', output_dir)
        group_resname = d.get('group_resname', group_resname)
        reference_ttl = d.get('reference_ttl', reference_ttl)
        min_stay      = d.get('min_stay', min_stay)
        min_gap       = d.get('min_gap', min_gap)
        max_gap       = d.get('max_gap', max_gap)
        band_dir      = d.get('band_dir', band_dir)
        separators    = d.get('separators', separators)
        starts        = d.get('starts', starts)

        print "# Using config from file: ", sys.argv[-1]
    else:
        print "# Using default (SAMPLE) config..."
        print "# Please review the config and make your own"
        print "# tailored to your needs."
        print "# To use your own file run: "
        print "#"
        print "#     python " + " ".join(sys.argv) + " my_file.conf"
        print "#"
        print "# my_file.conf should be a python source file"
        print "# (see help below). "

    in_conf = ""
    in_conf += "pdb_filename        = \"{0}\"\n".format(pdb_filename)
    in_conf += "xtc_filename        = \"{0}\"\n".format(xtc_filename)
    in_conf += "begin_frame         = {0}\n".format(begin_frame)
    in_conf += "end_frame           = {0}\n".format(end_frame)
    in_conf += "output_dir          = \"{0}\"\n".format(output_dir)
    in_conf += "group_resname       = \"{0}\"\n".format(group_resname)
    in_conf += "reference_ttl       = {0}\n".format(reference_ttl)
    in_conf += "min_stay            = {0}\n".format(min_stay)
    in_conf += "min_gap             = {0}\n".format(min_gap)
    in_conf += "max_gap             = {0}\n".format(max_gap)
    in_conf += "band_dir            = {0}\n".format(band_dir)
    in_conf += "separators          = {0}\n".format(separators)
    in_conf += "starts              = {0}\n".format(str(starts))
    in_conf += "end_frame           = {0}\n".format(end_frame)

    help_str = ""
    help_str += "###################################################\n"
    help_str += "# Command issued was:\n"
    help_str += "###################################################\n"
    help_str += "#     python " + " ".join(sys.argv) + "\n"
    help_str += "###################################################\n"
    help_str += "# Config file used:\n"
    help_str += "###################################################\n"
    help_str += "# (for meaning of variables see script file) ######\n"
    help_str += "###################################################\n"
    for item in in_conf.split('\n'):
        if item:  # no empty lines
            help_str += "# {0}\n".format(item)
    help_str += "###################################################\n"
    help_str += "# (input file ends here) ##########################\n"
    help_str += "# Remember to uncomment variables if want to use  #\n"
    help_str += "# them in a rerun. They are commented for gnuplot #\n"
    help_str += "# (with default gnuplot comment character '#'     #\n"
    help_str += "###################################################\n"

    # print help / summary
    print help_str

    reference_ttl = [10, 100]  # dla moich przykladow...

    # make output directory structure
    mkdir_p("{0}".format(output_dir))  # TODO: backup?

    # read data model
    data = Model(pdb_filename)

    # prepare useful indexes
    group_ndx_filename = "{0}/{1}.ndx".format(output_dir, group_resname)
    if not os.path.isfile(group_ndx_filename):
        for resname, ndx in indexes(data):
            ndx.write("{0}/{1}.ndx".format(output_dir, resname))

    # get the index for selected group resname
    # Notice: there is only one group in groups for this index!
    # we work on group to be faster
    selected_ndx = IndexFile(group_ndx_filename)

    bands = []
    left_bound = separators[0]
    for right_bound in separators[1:]:
        bands.append(Band(left_bound, right_bound, band_dir))
        left_bound = right_bound

    restarts = [[] for i in range(max(starts.keys())+1)]
    for key, r in starts.iteritems():
        restarts[key] = r
    # reload input data!
    initial_model = Model(pdb_filename)
    trajectory = Trajectory(xtc_filename)
    group = selected_ndx.groups[0].select_atoms(initial_model)

    counter = MSDSimplifiedDeviation(group, bands, restarts, reference_ttl)
    iterate(initial_model, trajectory, begin_frame, end_frame, counter.count)

    print "end_frame           = {0}\n".format(end_frame)

    print "AAAAAAAAAA"
    print counter.xs
    print counter.xf
    print counter.ys
    print counter.yf
    print counter.zs
    print counter.zf
    print "BBBBBBBBBB"

    for layer_no, band in enumerate(bands):
        # put to file with commments for future use
        out = file("{0}/deviation_layer_{1:03d}.csv".format(output_dir, layer_no), "w")
        xys = map(lambda a: sum(a), zip(counter.xs[layer_no], counter.ys[layer_no]))
        xyzs = map(lambda a: sum(a), zip(counter.xs[layer_no], counter.ys[layer_no], counter.zs[layer_no]))
        xyf = map(lambda a: sum(a), zip(counter.xf[layer_no], counter.yf[layer_no]))
        xyzf = map(lambda a: sum(a), zip(counter.xf[layer_no], counter.yf[layer_no], counter.zf[layer_no]))

        def simple_fit(t0, x0, t1, x1):
            return (x1 - x0) / float(t1 - t0)

        fit_fn = lambda x: simple_fit(reference_ttl[0], x[0], reference_ttl[1], x[1])
        xy = map(fit_fn, zip(xys, xyf))
        xyz = map(fit_fn, zip(xyzs, xyzf))
        x = map(fit_fn, zip(counter.xs[layer_no], counter.xf[layer_no]))
        y = map(fit_fn, zip(counter.ys[layer_no], counter.yf[layer_no]))
        z = map(fit_fn, zip(counter.zs[layer_no], counter.zf[layer_no]))

        xyz = [a * 1000.0 / 6.0 for a in xyz if a > 0]
        xy = [a * 1000.0 / 4.0 for a in xy if a > 0]
        x = [a * 1000.0 / 2.0 for a in x if a > 0]
        y = [a * 1000.0 / 2.0 for a in y if a > 0]
        z = [a * 1000.0 / 2.0 for a in z if a > 0]
        mlen = min(map(len, [xyz, xy, x, y, z]))
        if mlen:
            xyz = xyz[:mlen]
            xy = xy[:mlen]
            x = x[:mlen]
            y = y[:mlen]
            z = z[:mlen]
        else:
            xyz = []
            xy = []
            x = []
            y = []
            z = []

        def sddev(lst):
            if len(lst):
                av = sum(lst) / len(lst)
                var = sum(map(lambda x: (x - av) ** 2, lst)) / len(lst)
                return math.sqrt(var)
            else:
                return -1.0

        stdevs = map(sddev, [xyz, xy, x, y, z])

        data = zip(xyz, xy, x, y, z)
        # out.write(help_str)
        # out.write("#data in [nm^2]\n")
        out.write(" ".join(map(str, stdevs)) + "\n\n")
        out.write("#xyz xy x y z\n")
        for o in data:
            if all([x > 0 for x in o]):
                out.write(" ".join(map(str, o)) + "\n")

        out.close()

    print "DONE!"


# -*- coding: utf-8 -*-
"""
Sony Colourspaces
=================

Defines the *Sony* colourspaces:

-   :attr:`colour.models.S_GAMUT_COLOURSPACE`.
-   :attr:`colour.models.S_GAMUT3_COLOURSPACE`.
-   :attr:`colour.models.S_GAMUT3_CINE_COLOURSPACE`.
-   :attr:`colour.models.VENICE_S_GAMUT3_COLOURSPACE`.
-   :attr:`colour.models.VENICE_S_GAMUT3_CINE_COLOURSPACE`.

Notes
-----
-   The *Venice S-Gamut3* and *Venice S-Gamut3.Cine* primaries and whitepoint
    were derived with the following `Google Colab Notebook \
<https://colab.research.google.com/drive/1ZGTij7jT8eZRMPUkyWlv_x5ix5Q5twMB>`__.

References
----------
-   :cite:`Gaggioni` : Gaggioni, H., Dhanendra, P., Yamashita, J., Kawada, N.,
    Endo, K., & Clark, C. (n.d.). S-Log: A new LUT for digital production
    mastering and interchange applications (Vol. 709, pp. 1-13).
    http://pro.sony.com/bbsccms/assets/files/mkt/cinema/solutions/slog_manual.pdf
-   :cite:`SonyCorporation` : Sony Corporation. (n.d.). S-Log Whitepaper (pp.
    1-17). http://www.theodoropoulos.info/attachments/076_on%20S-Log.pdf
-   :cite:`SonyCorporationd` : Sony Corporation. (n.d.). Technical Summary
    for S-Gamut3.Cine/S-Log3 and S-Gamut3/S-Log3 (pp. 1-7).
    http://community.sony.com/sony/attachments/sony/\
large-sensor-camera-F5-F55/12359/2/\
TechnicalSummary_for_S-Gamut3Cine_S-Gamut3_S-Log3_V1_00.pdf
-   :cite:`SonyCorporatione` : Sony Corporation. (n.d.).
    S-Gamut3_S-Gamut3Cine_Matrix.xlsx.
    https://community.sony.com/sony/attachments/sony/\
large-sensor-camera-F5-F55/12359/3/S-Gamut3_S-Gamut3Cine_Matrix.xlsx
-   :cite:`SonyElectronicsCorporation2020` : Sony Electronics Corporation.
    (2020). IDT.Sony.Venice_SLog3_SGamut3.ctl. https://github.com/ampas/\
aces-dev/blob/710ecbe52c87ce9f4a1e02c8ddf7ea0d6b611cc8/transforms/ctl/idt/\
vendorSupplied/sony/IDT.Sony.Venice_SLog3_SGamut3.ctl
-   :cite:`SonyElectronicsCorporation2020a` : Sony Electronics Corporation.
    (2020). IDT.Sony.Venice_SLog3_SGamut3Cine.ctl. https://github.com/ampas/\
aces-dev/blob/710ecbe52c87ce9f4a1e02c8ddf7ea0d6b611cc8/transforms/ctl/idt/\
vendorSupplied/sony/IDT.Sony.Venice_SLog3_SGamut3Cine.ctl
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models.rgb import (RGB_Colourspace, log_encoding_SLog2,
                               log_decoding_SLog2, log_encoding_SLog3,
                               log_decoding_SLog3, normalised_primary_matrix)

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2020 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'S_GAMUT_PRIMARIES', 'S_GAMUT_WHITEPOINT_NAME', 'S_GAMUT_WHITEPOINT',
    'S_GAMUT_TO_XYZ_MATRIX', 'XYZ_TO_S_GAMUT_MATRIX', 'S_GAMUT_COLOURSPACE',
    'S_GAMUT3_PRIMARIES', 'S_GAMUT3_WHITEPOINT_NAME', 'S_GAMUT3_WHITEPOINT',
    'S_GAMUT3_TO_XYZ_MATRIX', 'XYZ_TO_S_GAMUT3_MATRIX', 'S_GAMUT3_COLOURSPACE',
    'S_GAMUT3_CINE_PRIMARIES', 'S_GAMUT3_CINE_WHITEPOINT_NAME',
    'S_GAMUT3_CINE_WHITEPOINT', 'S_GAMUT3_CINE_TO_XYZ_MATRIX',
    'XYZ_TO_S_GAMUT3_CINE_MATRIX', 'S_GAMUT3_CINE_COLOURSPACE',
    'VENICE_S_GAMUT3_PRIMARIES', 'VENICE_S_GAMUT3_WHITEPOINT_NAME',
    'VENICE_S_GAMUT3_WHITEPOINT', 'VENICE_S_GAMUT3_TO_XYZ_MATRIX',
    'XYZ_TO_VENICE_S_GAMUT3_MATRIX', 'VENICE_S_GAMUT3_COLOURSPACE',
    'VENICE_S_GAMUT3_CINE_PRIMARIES', 'VENICE_S_GAMUT3_CINE_WHITEPOINT_NAME',
    'VENICE_S_GAMUT3_CINE_WHITEPOINT', 'VENICE_S_GAMUT3_CINE_TO_XYZ_MATRIX',
    'XYZ_TO_VENICE_S_GAMUT3_CINE_MATRIX', 'VENICE_S_GAMUT3_CINE_COLOURSPACE'
]

S_GAMUT_PRIMARIES = np.array([
    [0.7300, 0.2800],
    [0.1400, 0.8550],
    [0.1000, -0.0500],
])
"""
*S-Gamut* colourspace primaries.

S_GAMUT_PRIMARIES : ndarray, (3, 2)
"""

S_GAMUT_WHITEPOINT_NAME = 'D65'
"""
*S-Gamut* colourspace whitepoint name.

S_GAMUT_WHITEPOINT_NAME : unicode
"""

S_GAMUT_WHITEPOINT = (ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
    S_GAMUT_WHITEPOINT_NAME])
"""
*S-Gamut* colourspace whitepoint.

S_GAMUT_WHITEPOINT : ndarray
"""

S_GAMUT_TO_XYZ_MATRIX = np.array([
    [0.7064827132, 0.1288010498, 0.1151721641],
    [0.2709796708, 0.7866064112, -0.0575860820],
    [-0.0096778454, 0.0046000375, 1.0941355587],
])
"""
*S-Gamut* colourspace to *CIE XYZ* tristimulus values matrix.

S_GAMUT_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_S_GAMUT_MATRIX = np.array([
    [1.5073998991, -0.2458221374, -0.1716116808],
    [-0.5181517271, 1.3553912409, 0.1258786682],
    [0.0155116982, -0.0078727714, 0.9119163656],
])
"""
*CIE XYZ* tristimulus values to *S-Gamut* colourspace matrix.

XYZ_TO_S_GAMUT_MATRIX : array_like, (3, 3)
"""

S_GAMUT_COLOURSPACE = RGB_Colourspace(
    'S-Gamut',
    S_GAMUT_PRIMARIES,
    S_GAMUT_WHITEPOINT,
    S_GAMUT_WHITEPOINT_NAME,
    S_GAMUT_TO_XYZ_MATRIX,
    XYZ_TO_S_GAMUT_MATRIX,
    log_encoding_SLog2,
    log_decoding_SLog2,
)
S_GAMUT_COLOURSPACE.__doc__ = """
*S-Gamut* colourspace.

References
----------
:cite:`Gaggioni`, :cite:`SonyCorporation`

S_GAMUT_COLOURSPACE : RGB_Colourspace
"""

S_GAMUT3_PRIMARIES = S_GAMUT_PRIMARIES
"""
*S-Gamut3* colourspace primaries.

S_GAMUT3_PRIMARIES : ndarray, (3, 2)
"""

S_GAMUT3_WHITEPOINT_NAME = S_GAMUT_WHITEPOINT_NAME
"""
*S-Gamut3* colourspace whitepoint name.

S_GAMUT3_WHITEPOINT_NAME : unicode
"""

S_GAMUT3_WHITEPOINT = S_GAMUT_WHITEPOINT
"""
*S-Gamut3* colourspace whitepoint.

S_GAMUT3_WHITEPOINT : ndarray
"""

S_GAMUT3_TO_XYZ_MATRIX = S_GAMUT_TO_XYZ_MATRIX
"""
*S-Gamut3* colourspace to *CIE XYZ* tristimulus values matrix.

S_GAMUT3_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_S_GAMUT3_MATRIX = XYZ_TO_S_GAMUT_MATRIX
"""
*CIE XYZ* tristimulus values to *S-Gamut3* colourspace matrix.

XYZ_TO_S_GAMUT3_MATRIX : array_like, (3, 3)
"""

S_GAMUT3_COLOURSPACE = RGB_Colourspace(
    'S-Gamut3',
    S_GAMUT3_PRIMARIES,
    S_GAMUT3_WHITEPOINT,
    S_GAMUT3_WHITEPOINT_NAME,
    S_GAMUT3_TO_XYZ_MATRIX,
    XYZ_TO_S_GAMUT3_MATRIX,
    log_encoding_SLog3,
    log_decoding_SLog3,
)
S_GAMUT3_COLOURSPACE.__doc__ = """
*S-Gamut3* colourspace.

References
----------
:cite:`SonyCorporationd`

S_GAMUT3_COLOURSPACE : RGB_Colourspace
"""

S_GAMUT3_CINE_PRIMARIES = np.array([
    [0.76600, 0.27500],
    [0.22500, 0.80000],
    [0.08900, -0.08700],
])
"""
*S-Gamut3.Cine* colourspace primaries.

S_GAMUT3_CINE_PRIMARIES : ndarray, (3, 2)
"""

S_GAMUT3_CINE_WHITEPOINT_NAME = S_GAMUT_WHITEPOINT_NAME
"""
*S-Gamut3.Cine* colourspace whitepoint name.

S_GAMUT3_CINE_WHITEPOINT_NAME : unicode
"""

S_GAMUT3_CINE_WHITEPOINT = S_GAMUT_WHITEPOINT
"""
*S-Gamut3.Cine* colourspace whitepoint.

S_GAMUT3_CINE_WHITEPOINT : ndarray
"""

S_GAMUT3_CINE_TO_XYZ_MATRIX = np.array([
    [0.5990839208, 0.2489255161, 0.1024464902],
    [0.2150758201, 0.8850685017, -0.1001443219],
    [-0.0320658495, -0.0276583907, 1.1487819910],
])
"""
*S-Gamut3.Cine* colourspace to *CIE XYZ* tristimulus values matrix.

S_GAMUT3_CINE_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_S_GAMUT3_CINE_MATRIX = np.array([
    [1.8467789693, -0.5259861230, -0.2105452114],
    [-0.4441532629, 1.2594429028, 0.1493999729],
    [0.0408554212, 0.0156408893, 0.8682072487],
])
"""
*CIE XYZ* tristimulus values to *S-Gamut3.Cine* colourspace matrix.

XYZ_TO_S_GAMUT3_CINE_MATRIX : array_like, (3, 3)
"""

S_GAMUT3_CINE_COLOURSPACE = RGB_Colourspace(
    'S-Gamut3.Cine',
    S_GAMUT3_CINE_PRIMARIES,
    S_GAMUT3_CINE_WHITEPOINT,
    S_GAMUT3_CINE_WHITEPOINT_NAME,
    S_GAMUT3_CINE_TO_XYZ_MATRIX,
    XYZ_TO_S_GAMUT3_CINE_MATRIX,
    log_encoding_SLog3,
    log_decoding_SLog3,
)
S_GAMUT3_CINE_COLOURSPACE.__doc__ = """
*S-Gamut3.Cine* colourspace.

References
----------
:cite:`SonyCorporatione`

S_GAMUT3_CINE_COLOURSPACE : RGB_Colourspace
"""

VENICE_S_GAMUT3_PRIMARIES = np.array([
    [0.740464264304292, 0.279364374750660],
    [0.089241145423286, 0.893809528608105],
    [0.110488236673827, -0.052579333080476],
])
"""
*Venice S-Gamut3* colourspace primaries.

VENICE_S_GAMUT3_PRIMARIES : ndarray, (3, 2)
"""

VENICE_S_GAMUT3_WHITEPOINT_NAME = S_GAMUT_WHITEPOINT_NAME
"""
*Venice S-Gamut3* colourspace whitepoint name.

VENICE_S_GAMUT3_WHITEPOINT_NAME : unicode
"""

VENICE_S_GAMUT3_WHITEPOINT = S_GAMUT_WHITEPOINT
"""
*Venice S-Gamut3* colourspace whitepoint.

VENICE_S_GAMUT3_WHITEPOINT : ndarray
"""

VENICE_S_GAMUT3_TO_XYZ_MATRIX = normalised_primary_matrix(
    VENICE_S_GAMUT3_PRIMARIES, VENICE_S_GAMUT3_WHITEPOINT)
"""
*Venice S-Gamut3* colourspace to *CIE XYZ* tristimulus values matrix.

VENICE_S_GAMUT3_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_VENICE_S_GAMUT3_MATRIX = np.linalg.inv(VENICE_S_GAMUT3_TO_XYZ_MATRIX)
"""
*CIE XYZ* tristimulus values to *Venice S-Gamut3* colourspace matrix.

XYZ_TO_VENICE_S_GAMUT3_MATRIX : array_like, (3, 3)
"""

VENICE_S_GAMUT3_COLOURSPACE = RGB_Colourspace(
    'Venice S-Gamut3',
    VENICE_S_GAMUT3_PRIMARIES,
    VENICE_S_GAMUT3_WHITEPOINT,
    VENICE_S_GAMUT3_WHITEPOINT_NAME,
    VENICE_S_GAMUT3_TO_XYZ_MATRIX,
    XYZ_TO_VENICE_S_GAMUT3_MATRIX,
    log_encoding_SLog3,
    log_decoding_SLog3,
)
VENICE_S_GAMUT3_COLOURSPACE.__doc__ = """
*Venice S-Gamut3* colourspace.

References
----------
:cite:`SonyElectronicsCorporation2020`

VENICE_S_GAMUT3_COLOURSPACE : RGB_Colourspace
"""

VENICE_S_GAMUT3_CINE_PRIMARIES = np.array([
    [0.775901871567345, 0.274502392854799],
    [0.188682902773355, 0.828684937020288],
    [0.101337382499301, -0.089187517306263],
])
"""
*Venice S-Gamut3.Cine* colourspace primaries.

VENICE_S_GAMUT3_CINE_PRIMARIES : ndarray, (3, 2)
"""

VENICE_S_GAMUT3_CINE_WHITEPOINT_NAME = S_GAMUT_WHITEPOINT_NAME
"""
*Venice S-Gamut3.Cine* colourspace whitepoint name.

VENICE_S_GAMUT3_CINE_WHITEPOINT_NAME : unicode
"""

VENICE_S_GAMUT3_CINE_WHITEPOINT = S_GAMUT_WHITEPOINT
"""
*Venice S-Gamut3.Cine* colourspace whitepoint.

VENICE_S_GAMUT3_CINE_WHITEPOINT : ndarray
"""

VENICE_S_GAMUT3_CINE_TO_XYZ_MATRIX = normalised_primary_matrix(
    VENICE_S_GAMUT3_CINE_PRIMARIES, VENICE_S_GAMUT3_CINE_WHITEPOINT)
"""
*Venice S-Gamut3.Cine* colourspace to *CIE XYZ* tristimulus values matrix.

VENICE_S_GAMUT3_CINE_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_VENICE_S_GAMUT3_CINE_MATRIX = np.linalg.inv(
    VENICE_S_GAMUT3_CINE_TO_XYZ_MATRIX)
"""
*CIE XYZ* tristimulus values to *Venice S-Gamut3.Cine* colourspace matrix.

XYZ_TO_VENICE_S_GAMUT3_CINE_MATRIX : array_like, (3, 3)
"""

VENICE_S_GAMUT3_CINE_COLOURSPACE = RGB_Colourspace(
    'Venice S-Gamut3.Cine',
    VENICE_S_GAMUT3_CINE_PRIMARIES,
    VENICE_S_GAMUT3_CINE_WHITEPOINT,
    VENICE_S_GAMUT3_CINE_WHITEPOINT_NAME,
    VENICE_S_GAMUT3_CINE_TO_XYZ_MATRIX,
    XYZ_TO_VENICE_S_GAMUT3_CINE_MATRIX,
    log_encoding_SLog3,
    log_decoding_SLog3,
)
VENICE_S_GAMUT3_CINE_COLOURSPACE.__doc__ = """
*Venice S-Gamut3.Cine* colourspace.

References
----------
:cite:`SonyElectronicsCorporation2020a`

VENICE_S_GAMUT3_CINE_COLOURSPACE : RGB_Colourspace
"""


#
# Implements the core metrics from sound event detection evaluation module http://tut-arg.github.io/sed_eval/ and
# The DOA metrics are explained in the SELDnet paper
#
# This script has MIT license
#

import numpy as np
from scipy.optimize import linear_sum_assignment
eps = np.finfo(np.float).eps

##########################################################################################
# SELD scoring functions - class implementation
#
# NOTE: Supports only one-hot labels for both SED and DOA. Doesnt work for baseline method
# directly, since it estimated DOA in regression approach. Check below the class for
# one shot (function) implementations of all metrics. The function implementation has
# support for both one-hot labels and regression values of DOA estimation.
##########################################################################################

class SELDMetrics(object):
    def __init__(self, nb_frames_1s=None, data_gen=None):
        # SED params
        self._S = 0
        self._D = 0
        self._I = 0
        self._TP = 0
        self._Nref = 0
        self._Nsys = 0
        self._block_size = nb_frames_1s

        # DOA params
        self._doa_loss_pred_cnt = 0
        self._nb_frames = 0

        self._doa_loss_pred = 0
        self._nb_good_pks = 0

        self._data_gen = data_gen

        self._less_est_cnt, self._less_est_frame_cnt = 0, 0
        self._more_est_cnt, self._more_est_frame_cnt = 0, 0

    def f1_overall_framewise(self, O, T):
        TP = ((2 * T - O) == 1).sum()
        Nref, Nsys = T.sum(), O.sum()
        self._TP += TP
        self._Nref += Nref
        self._Nsys += Nsys

    def er_overall_framewise(self, O, T):
        FP = np.logical_and(T == 0, O == 1).sum(1)
        FN = np.logical_and(T == 1, O == 0).sum(1)
        S = np.minimum(FP, FN).sum()
        D = np.maximum(0, FN - FP).sum()
        I = np.maximum(0, FP - FN).sum()
        self._S += S
        self._D += D
        self._I += I

    def f1_overall_1sec(self, O, T):
        new_size = int(np.ceil(float(O.shape[0]) / self._block_size))
        O_block = np.zeros((new_size, O.shape[1]))
        T_block = np.zeros((new_size, O.shape[1]))
        for i in range(0, new_size):
            O_block[i, :] = np.max(O[int(i * self._block_size):int(i * self._block_size + self._block_size - 1), :], axis=0)
            T_block[i, :] = np.max(T[int(i * self._block_size):int(i * self._block_size + self._block_size - 1), :], axis=0)
        return self.f1_overall_framewise(O_block, T_block)

    def er_overall_1sec(self, O, T):
        new_size = int(np.ceil(float(O.shape[0]) / self._block_size))
        O_block = np.zeros((new_size, O.shape[1]))
        T_block = np.zeros((new_size, O.shape[1]))
        for i in range(0, new_size):
            O_block[i, :] = np.max(O[int(i * self._block_size):int(i * self._block_size + self._block_size - 1), :], axis=0)
            T_block[i, :] = np.max(T[int(i * self._block_size):int(i * self._block_size + self._block_size - 1), :], axis=0)
        return self.er_overall_framewise(O_block, T_block)

    def update_sed_scores(self, pred, gt):
        """
        Computes SED metrics for one second segments

        :param pred: predicted matrix of dimension [nb_frames, nb_classes], with 1 when sound event is active else 0
        :param gt:  reference matrix of dimension [nb_frames, nb_classes], with 1 when sound event is active else 0
        :param nb_frames_1s: integer, number of frames in one second
        :return:
        """
        self.f1_overall_1sec(pred, gt)
        self.er_overall_1sec(pred, gt)

    def compute_sed_scores(self):
        ER = (self._S + self._D + self._I) / (self._Nref + 0.0)

        prec = float(self._TP) / float(self._Nsys + eps)
        recall = float(self._TP) / float(self._Nref + eps)
        F = 2 * prec * recall / (prec + recall + eps)

        return ER, F

    def update_doa_scores(self, pred_doa_thresholded, gt_doa):
        '''
        Compute DOA metrics when DOA is estimated using classification approach

        :param pred_doa_thresholded: predicted results of dimension [nb_frames, nb_classes, nb_azi*nb_ele],
                                    with value 1 when sound event active, else 0
        :param gt_doa: reference results of dimension [nb_frames, nb_classes, nb_azi*nb_ele],
                        with value 1 when sound event active, else 0
        :param data_gen_test: feature or data generator class

        :return: DOA metrics

        '''
        self._doa_loss_pred_cnt += np.sum(pred_doa_thresholded)
        self._nb_frames += pred_doa_thresholded.shape[0]

        for frame in range(pred_doa_thresholded.shape[0]):
            nb_gt_peaks = int(np.sum(gt_doa[frame, :]))
            nb_pred_peaks = int(np.sum(pred_doa_thresholded[frame, :]))

            # good_frame_cnt includes frames where the nb active sources were zero in both groundtruth and prediction
            if nb_gt_peaks == nb_pred_peaks:
                self._nb_good_pks += 1
            elif nb_gt_peaks > nb_pred_peaks:
                self._less_est_frame_cnt += 1
                self._less_est_cnt += (nb_gt_peaks - nb_pred_peaks)
            elif nb_pred_peaks > nb_gt_peaks:
                self._more_est_frame_cnt += 1
                self._more_est_cnt += (nb_pred_peaks - nb_gt_peaks)

            # when nb_ref_doa > nb_estimated_doa, ignores the extra ref doas and scores only the nearest matching doas
            # similarly, when nb_estimated_doa > nb_ref_doa, ignores the extra estimated doa and scores the remaining matching doas
            if nb_gt_peaks and nb_pred_peaks:
                pred_ind = np.where(pred_doa_thresholded[frame] == 1)[1]
                pred_list_rad = np.array(self._data_gen .get_matrix_index(pred_ind)) * np.pi / 180

                gt_ind = np.where(gt_doa[frame] == 1)[1]
                gt_list_rad = np.array(self._data_gen .get_matrix_index(gt_ind)) * np.pi / 180

                frame_dist = distance_between_gt_pred(gt_list_rad.T, pred_list_rad.T)
                self._doa_loss_pred += frame_dist

    def compute_doa_scores(self):
        doa_error = self._doa_loss_pred / self._doa_loss_pred_cnt
        frame_recall = self._nb_good_pks / float(self._nb_frames)
        return doa_error, frame_recall

    def reset(self):
        # SED params
        self._S = 0
        self._D = 0
        self._I = 0
        self._TP = 0
        self._Nref = 0
        self._Nsys = 0

        # DOA params
        self._doa_loss_pred_cnt = 0
        self._nb_frames = 0

        self._doa_loss_pred = 0
        self._nb_good_pks = 0

        self._less_est_cnt, self._less_est_frame_cnt = 0, 0
        self._more_est_cnt, self._more_est_frame_cnt = 0, 0

###############################################################
# SED scoring functions
###############################################################

def reshape_3Dto2D(A):
    return A.reshape(A.shape[0] * A.shape[1], A.shape[2])

def f1_overall_framewise(O, T):
    if len(O.shape) == 3:
        O, T = reshape_3Dto2D(O), reshape_3Dto2D(T)
    TP = ((2 * T - O) == 1).sum()
    Nref, Nsys = T.sum(), O.sum()

    prec = float(TP) / float(Nsys + eps)
    recall = float(TP) / float(Nref + eps)
    f1_score = 2 * prec * recall / (prec + recall + eps)
    return f1_score

def er_overall_framewise(O, T):
    if len(O.shape) == 3:
        O, T = reshape_3Dto2D(O), reshape_3Dto2D(T)

    FP = np.logical_and(T == 0, O == 1).sum(1)
    FN = np.logical_and(T == 1, O == 0).sum(1)

    S = np.minimum(FP, FN).sum()
    D = np.maximum(0, FN-FP).sum()
    I = np.maximum(0, FP-FN).sum()

    Nref = T.sum()
    ER = (S+D+I) / (Nref + 0.0)
    return ER

def f1_overall_1sec(O, T, block_size):
    if len(O.shape) == 3:
        O, T = reshape_3Dto2D(O), reshape_3Dto2D(T)
    new_size = int(np.ceil(float(O.shape[0]) / block_size))
    O_block = np.zeros((new_size, O.shape[1]))
    T_block = np.zeros((new_size, O.shape[1]))
    for i in range(0, new_size):
        O_block[i, :] = np.max(O[int(i * block_size):int(i * block_size + block_size - 1), :], axis=0)
        T_block[i, :] = np.max(T[int(i * block_size):int(i * block_size + block_size - 1), :], axis=0)
    return f1_overall_framewise(O_block, T_block)

def er_overall_1sec(O, T, block_size):
    if len(O.shape) == 3:
        O, T = reshape_3Dto2D(O), reshape_3Dto2D(T)
    new_size = int(np.ceil(float(O.shape[0]) / block_size))
    O_block = np.zeros((new_size, O.shape[1]))
    T_block = np.zeros((new_size, O.shape[1]))
    for i in range(0, new_size):
        O_block[i, :] = np.max(O[int(i * block_size):int(i * block_size + block_size - 1), :], axis=0)
        T_block[i, :] = np.max(T[int(i * block_size):int(i * block_size + block_size - 1), :], axis=0)
    return er_overall_framewise(O_block, T_block)

def compute_sed_scores(pred, gt, nb_frames_1s):
    """
    Computes SED metrics for one second segments

    :param pred: predicted matrix of dimension [nb_frames, nb_classes], with 1 when sound event is active else 0
    :param gt:  reference matrix of dimension [nb_frames, nb_classes], with 1 when sound event is active else 0
    :param nb_frames_1s: integer, number of frames in one second
    :return:
    """
    f1o = f1_overall_1sec(pred, gt, nb_frames_1s)
    ero = er_overall_1sec(pred, gt, nb_frames_1s)
    scores = [ero, f1o]
    return scores

###############################################################
# DOA scoring functions
###############################################################

def compute_doa_scores_regr_xyz(pred_doa, gt_doa, pred_sed, gt_sed):
    """
        Compute DOA metrics when DOA is estimated using regression approach

    :param pred_doa: predicted doa_labels is of dimension [nb_frames, 3*nb_classes],
                        nb_classes each for x, y, and z axes,
                        if active, the DOA values will be in real numbers [-1 1] range, else, it will contain default doa values of (0, 0, 0)
    :param gt_doa: reference doa_labels is of dimension [nb_frames, 3*nb_classes],
    :param pred_sed: predicted sed label of dimension [nb_frames, nb_classes] which is 1 for active sound event else zero
    :param gt_sed: reference sed label of dimension [nb_frames, nb_classes] which is 1 for active sound event else zero
    :return:
    """

    nb_src_gt_list = np.zeros(gt_doa.shape[0]).astype(int)
    nb_src_pred_list = np.zeros(gt_doa.shape[0]).astype(int)
    good_frame_cnt = 0
    doa_loss_pred = 0.0
    nb_sed = gt_sed.shape[-1]

    less_est_cnt, less_est_frame_cnt = 0, 0
    more_est_cnt, more_est_frame_cnt = 0, 0

    for frame_cnt, sed_frame in enumerate(gt_sed):
        nb_src_gt_list[frame_cnt] = int(np.sum(sed_frame))
        nb_src_pred_list[frame_cnt] = int(np.sum(pred_sed[frame_cnt]))

        # good_frame_cnt includes frames where the nb active sources were zero in both groundtruth and prediction
        if nb_src_gt_list[frame_cnt] == nb_src_pred_list[frame_cnt]:
            good_frame_cnt = good_frame_cnt + 1
        elif nb_src_gt_list[frame_cnt] > nb_src_pred_list[frame_cnt]:
            less_est_cnt = less_est_cnt + nb_src_gt_list[frame_cnt] - nb_src_pred_list[frame_cnt]
            less_est_frame_cnt = less_est_frame_cnt + 1
        elif nb_src_gt_list[frame_cnt] < nb_src_pred_list[frame_cnt]:
            more_est_cnt = more_est_cnt + nb_src_pred_list[frame_cnt] - nb_src_gt_list[frame_cnt]
            more_est_frame_cnt = more_est_frame_cnt + 1

        # when nb_ref_doa > nb_estimated_doa, ignores the extra ref doas and scores only the nearest matching doas
        # similarly, when nb_estimated_doa > nb_ref_doa, ignores the extra estimated doa and scores the remaining matching doas
        if nb_src_gt_list[frame_cnt] and nb_src_pred_list[frame_cnt]:
            # DOA Loss with respect to predicted confidence
            sed_frame_gt = gt_sed[frame_cnt]
            doa_frame_gt_x = gt_doa[frame_cnt][:nb_sed][sed_frame_gt == 1]
            doa_frame_gt_y = gt_doa[frame_cnt][nb_sed:2*nb_sed][sed_frame_gt == 1]
            doa_frame_gt_z = gt_doa[frame_cnt][2*nb_sed:][sed_frame_gt == 1]

            sed_frame_pred = pred_sed[frame_cnt]
            doa_frame_pred_x = pred_doa[frame_cnt][:nb_sed][sed_frame_pred == 1]
            doa_frame_pred_y = pred_doa[frame_cnt][nb_sed:2*nb_sed][sed_frame_pred == 1]
            doa_frame_pred_z = pred_doa[frame_cnt][2*nb_sed:][sed_frame_pred == 1]

            doa_loss_pred += distance_between_gt_pred_xyz(np.vstack((doa_frame_gt_x, doa_frame_gt_y, doa_frame_gt_z)).T,
                                                      np.vstack((doa_frame_pred_x, doa_frame_pred_y, doa_frame_pred_z)).T)

    doa_loss_pred_cnt = np.sum(nb_src_pred_list)
    if doa_loss_pred_cnt:
        doa_loss_pred /= doa_loss_pred_cnt

    frame_recall = good_frame_cnt / float(gt_sed.shape[0])
    er_metric = [doa_loss_pred, frame_recall, doa_loss_pred_cnt, good_frame_cnt, more_est_cnt, less_est_cnt]
    return er_metric

def compute_doa_scores_regr(pred_doa_rad, gt_doa_rad, pred_sed, gt_sed):
    """
        Compute DOA metrics when DOA is estimated using regression approach

    :param pred_doa_rad: predicted doa_labels is of dimension [nb_frames, 2*nb_classes],
                        nb_classes each for azimuth and elevation angles,
                        if active, the DOA values will be in RADIANS, else, it will contain default doa values
    :param gt_doa_rad: reference doa_labels is of dimension [nb_frames, 2*nb_classes],
                    nb_classes each for azimuth and elevation angles,
                    if active, the DOA values will be in RADIANS, else, it will contain default doa values
    :param pred_sed: predicted sed label of dimension [nb_frames, nb_classes] which is 1 for active sound event else zero
    :param gt_sed: reference sed label of dimension [nb_frames, nb_classes] which is 1 for active sound event else zero
    :return:
    """

    nb_src_gt_list = np.zeros(gt_doa_rad.shape[0]).astype(int)
    nb_src_pred_list = np.zeros(gt_doa_rad.shape[0]).astype(int)
    good_frame_cnt = 0
    doa_loss_pred = 0.0
    nb_sed = gt_sed.shape[-1]

    less_est_cnt, less_est_frame_cnt = 0, 0
    more_est_cnt, more_est_frame_cnt = 0, 0

    for frame_cnt, sed_frame in enumerate(gt_sed):
        nb_src_gt_list[frame_cnt] = int(np.sum(sed_frame))
        nb_src_pred_list[frame_cnt] = int(np.sum(pred_sed[frame_cnt]))

        # good_frame_cnt includes frames where the nb active sources were zero in both groundtruth and prediction
        if nb_src_gt_list[frame_cnt] == nb_src_pred_list[frame_cnt]:
            good_frame_cnt = good_frame_cnt + 1
        elif nb_src_gt_list[frame_cnt] > nb_src_pred_list[frame_cnt]:
            less_est_cnt = less_est_cnt + nb_src_gt_list[frame_cnt] - nb_src_pred_list[frame_cnt]
            less_est_frame_cnt = less_est_frame_cnt + 1
        elif nb_src_gt_list[frame_cnt] < nb_src_pred_list[frame_cnt]:
            more_est_cnt = more_est_cnt + nb_src_pred_list[frame_cnt] - nb_src_gt_list[frame_cnt]
            more_est_frame_cnt = more_est_frame_cnt + 1

        # when nb_ref_doa > nb_estimated_doa, ignores the extra ref doas and scores only the nearest matching doas
        # similarly, when nb_estimated_doa > nb_ref_doa, ignores the extra estimated doa and scores the remaining matching doas
        if nb_src_gt_list[frame_cnt] and nb_src_pred_list[frame_cnt]:
            # DOA Loss with respect to predicted confidence
            sed_frame_gt = gt_sed[frame_cnt]
            doa_frame_gt_azi = gt_doa_rad[frame_cnt][:nb_sed][sed_frame_gt == 1]
            doa_frame_gt_ele = gt_doa_rad[frame_cnt][nb_sed:][sed_frame_gt == 1]

            sed_frame_pred = pred_sed[frame_cnt]
            doa_frame_pred_azi = pred_doa_rad[frame_cnt][:nb_sed][sed_frame_pred == 1]
            doa_frame_pred_ele = pred_doa_rad[frame_cnt][nb_sed:][sed_frame_pred == 1]

            doa_loss_pred += distance_between_gt_pred(np.vstack((doa_frame_gt_azi, doa_frame_gt_ele)).T,
                                                      np.vstack((doa_frame_pred_azi, doa_frame_pred_ele)).T)

    doa_loss_pred_cnt = np.sum(nb_src_pred_list)
    if doa_loss_pred_cnt:
        doa_loss_pred /= doa_loss_pred_cnt

    frame_recall = good_frame_cnt / float(gt_sed.shape[0])
    er_metric = [doa_loss_pred, frame_recall, doa_loss_pred_cnt, good_frame_cnt, more_est_cnt, less_est_cnt]
    return er_metric

def compute_doa_scores_clas(pred_doa_thresholded, gt_doa, data_gen_test):
    '''
    Compute DOA metrics when DOA is estimated using classification approach

    :param pred_doa_thresholded: predicted results of dimension [nb_frames, nb_classes, nb_azi*nb_ele],
                                with value 1 when sound event active, else 0
    :param gt_doa: reference results of dimension [nb_frames, nb_classes, nb_azi*nb_ele],
                    with value 1 when sound event active, else 0
    :param data_gen_test: feature or data generator class

    :return: DOA metrics

    '''
    doa_loss_pred_cnt = np.sum(pred_doa_thresholded)

    doa_loss_pred = 0
    nb_good_pks = 0

    less_est_cnt, less_est_frame_cnt = 0, 0
    more_est_cnt, more_est_frame_cnt = 0, 0

    for frame in range(pred_doa_thresholded.shape[0]):
        nb_gt_peaks = int(np.sum(gt_doa[frame, :]))
        nb_pred_peaks = int(np.sum(pred_doa_thresholded[frame, :]))

        # good_frame_cnt includes frames where the nb active sources were zero in both groundtruth and prediction
        if nb_gt_peaks == nb_pred_peaks:
            nb_good_pks += 1
        elif nb_gt_peaks > nb_pred_peaks:
            less_est_frame_cnt += 1
            less_est_cnt += (nb_gt_peaks - nb_pred_peaks)
        elif nb_pred_peaks > nb_gt_peaks:
            more_est_frame_cnt += 1
            more_est_cnt += (nb_pred_peaks - nb_gt_peaks)

        # when nb_ref_doa > nb_estimated_doa, ignores the extra ref doas and scores only the nearest matching doas
        # similarly, when nb_estimated_doa > nb_ref_doa, ignores the extra estimated doa and scores the remaining matching doas
        if nb_gt_peaks and nb_pred_peaks:
            pred_ind = np.where(pred_doa_thresholded[frame] == 1)[1]
            pred_list_rad = np.array(data_gen_test.get_matrix_index(pred_ind)) * np.pi / 180

            gt_ind = np.where(gt_doa[frame] == 1)[1]
            gt_list_rad = np.array(data_gen_test.get_matrix_index(gt_ind)) * np.pi / 180

            frame_dist = distance_between_gt_pred(gt_list_rad.T, pred_list_rad.T)
            doa_loss_pred += frame_dist

    if doa_loss_pred_cnt:
        doa_loss_pred /= doa_loss_pred_cnt

    frame_recall = nb_good_pks / float(pred_doa_thresholded.shape[0])
    er_metric = [doa_loss_pred, frame_recall, doa_loss_pred_cnt, nb_good_pks, more_est_cnt, less_est_cnt]
    return er_metric

def distance_between_gt_pred(gt_list_rad, pred_list_rad):
    """
    Shortest distance between two sets of spherical coordinates. Given a set of groundtruth spherical coordinates,
     and its respective predicted coordinates, we calculate the spherical distance between each of the spherical
     coordinate pairs resulting in a matrix of distances, where one axis represents the number of groundtruth
     coordinates and the other the predicted coordinates. The number of estimated peaks need not be the same as in
     groundtruth, thus the distance matrix is not always a square matrix. We use the hungarian algorithm to find the
     least cost in this distance matrix.

    :param gt_list_rad: list of ground-truth spherical coordinates
    :param pred_list_rad: list of predicted spherical coordinates
    :return: cost -  distance
    :return: less - number of DOA's missed
    :return: extra - number of DOA's over-estimated
    """

    gt_len, pred_len = gt_list_rad.shape[0], pred_list_rad.shape[0]
    ind_pairs = np.array([[x, y] for y in range(pred_len) for x in range(gt_len)])
    cost_mat = np.zeros((gt_len, pred_len))

    # Slow implementation
    # cost_mat = np.zeros((gt_len, pred_len))
    # for gt_cnt, gt in enumerate(gt_list_rad):
    #     for pred_cnt, pred in enumerate(pred_list_rad):
    #         cost_mat[gt_cnt, pred_cnt] = distance_between_spherical_coordinates_rad(gt, pred)

    # Fast implementation
    if gt_len and pred_len:
        az1, ele1, az2, ele2 = gt_list_rad[ind_pairs[:, 0], 0], gt_list_rad[ind_pairs[:, 0], 1], \
                               pred_list_rad[ind_pairs[:, 1], 0], pred_list_rad[ind_pairs[:, 1], 1]
        cost_mat[ind_pairs[:, 0], ind_pairs[:, 1]] = distance_between_spherical_coordinates_rad(az1, ele1, az2, ele2)

    row_ind, col_ind = linear_sum_assignment(cost_mat)
    cost = cost_mat[row_ind, col_ind].sum()
    return cost

def distance_between_gt_pred_xyz(gt_list, pred_list):
    """
    Shortest distance between two sets of Cartesian coordinates. Given a set of groundtruth coordinates,
     and its respective predicted coordinates, we calculate the spherical distance between each of the spherical
     coordinate pairs resulting in a matrix of distances, where one axis represents the number of groundtruth
     coordinates and the other the predicted coordinates. The number of estimated peaks need not be the same as in
     groundtruth, thus the distance matrix is not always a square matrix. We use the hungarian algorithm to find the
     least cost in this distance matrix.

    :param gt_list: list of ground-truth Cartesian coordinates
    :param pred_list: list of predicted Cartesian coordinates
    :return: cost -  distance
    :return: less - number of DOA's missed
    :return: extra - number of DOA's over-estimated
    """

    gt_len, pred_len = gt_list.shape[0], pred_list.shape[0]
    ind_pairs = np.array([[x, y] for y in range(pred_len) for x in range(gt_len)])
    cost_mat = np.zeros((gt_len, pred_len))

    # Slow implementation
    # cost_mat = np.zeros((gt_len, pred_len))
    # for gt_cnt, gt in enumerate(gt_list_rad):
    #     for pred_cnt, pred in enumerate(pred_list_rad):
    #         cost_mat[gt_cnt, pred_cnt] = distance_between_spherical_coordinates_rad(gt, pred)

    # Fast implementation
    if gt_len and pred_len:
        x1, y1, z1, x2, y2, z2 = gt_list[ind_pairs[:, 0], 0], gt_list[ind_pairs[:, 0], 1], gt_list[ind_pairs[:, 0], 2], \
                               pred_list[ind_pairs[:, 1], 0], pred_list[ind_pairs[:, 1], 1], pred_list[ind_pairs[:, 1], 2]
        cost_mat[ind_pairs[:, 0], ind_pairs[:, 1]] = distance_between_cartesian_coordinates(x1, y1, z1, x2, y2, z2)

    row_ind, col_ind = linear_sum_assignment(cost_mat)
    cost = cost_mat[row_ind, col_ind].sum()
    return cost

def distance_between_spherical_coordinates_rad(az1, ele1, az2, ele2):
    """
    Angular distance between two spherical coordinates
    MORE: https://en.wikipedia.org/wiki/Great-circle_distance

    :return: angular distance in degrees
    """
    dist = np.sin(ele1) * np.sin(ele2) + np.cos(ele1) * np.cos(ele2) * np.cos(np.abs(az1 - az2))
    # Making sure the dist values are in -1 to 1 range, else np.arccos kills the job
    dist = np.clip(dist, -1, 1)
    dist = np.arccos(dist) * 180 / np.pi
    return dist

def distance_between_cartesian_coordinates(x1, y1, z1, x2, y2, z2):
    """
    Angular distance between two cartesian coordinates
    MORE: https://en.wikipedia.org/wiki/Great-circle_distance
    Check 'From chord length' section

    :return: angular distance in degrees
    """
    # Normalize the Cartesian vectors
    N1 = np.sqrt(x1**2 + y1**2 + z1**2 + 1e-10)
    N2 = np.sqrt(x2**2 + y2**2 + z2**2 + 1e-10)
    x1, y1, z1, x2, y2, z2 = x1/N1, y1/N1, z1/N1, x2/N2, y2/N2, z2/N2

    #Compute the distance
    dist = x1*x2 + y1*y2 + z1*z2
    dist = np.clip(dist, -1, 1)
    dist = np.arccos(dist) * 180 / np.pi
    return dist

def sph2cart(azimuth, elevation, r):
    '''
    Convert spherical to cartesian coordinates

    :param azimuth: in radians
    :param elevation: in radians
    :param r: in meters
    :return: cartesian coordinates
    '''

    x = r * np.cos(elevation) * np.cos(azimuth)
    y = r * np.cos(elevation) * np.sin(azimuth)
    z = r * np.sin(elevation)
    return x, y, z

def cart2sph(x, y, z):
    '''
    Convert cartesian to spherical coordinates

    :param x:
    :param y:
    :param z:
    :return: azi, ele in radians and r in meters
    '''

    azimuth = np.arctan2(y,x)
    elevation = np.arctan2(z,np.sqrt(x**2 + y**2))
    r = np.sqrt(x**2 + y**2 + z**2)
    return azimuth, elevation, r

###############################################################
# SELD scoring functions
###############################################################

def early_stopping_metric(sed_error, doa_error):
    """
    Compute early stopping metric from sed and doa errors.

    :param sed_error: [error rate (0 to 1 range), f score (0 to 1 range)]
    :param doa_error: [doa error (in degrees), frame recall (0 to 1 range)]
    :return: seld metric result
    """
    seld_metric = np.mean([
        sed_error[0],
        1 - sed_error[1],
        doa_error[0]/180,
        1 - doa_error[1]]
        )
    return seld_metric


"""
Generate a sample using the Method of Morris

Three variants of Morris' sampling for elementary effects is supported:

- Vanilla Morris
- Optimised trajectories when ``optimal_trajectories=True`` (using
    Campolongo's enhancements from 2007 and optionally Ruano's enhancement
    from 2012; ``local_optimization=True``)
- Groups with optimised trajectories when ``optimal_trajectories=True`` and
    the problem definition specifies groups (note that ``local_optimization``
    must be ``False``)

At present, optimised trajectories is implemented using either a brute-force
approach, which can be very slow, especially if you require more than four
trajectories, or a local method based which is much faster. Both methods now
implement working with groups of factors.

Note that the number of factors makes little difference,
but the ratio between number of optimal trajectories and the sample size
results in an exponentially increasing number of scores that must be
computed to find the optimal combination of trajectories.  We suggest going
no higher than 4 from a pool of 100 samples with the brute force approach.
With local_optimization = True (which is default),
it is possible to go higher than the previously suggested 4 from 100.

"""
from __future__ import division

import numpy as np

import numpy.random as rd

from . gurobi import GlobalOptimisation
from . local import LocalOptimisation
from . brute import BruteForce

from . strategy import SampleMorris

from SALib.sample import common_args
from SALib.util import scale_samples, read_param_file, compute_groups_matrix

try:
    import gurobipy
except ImportError:
    _has_gurobi = False
else:
    _has_gurobi = True

__all__ = ['sample']

def sample(problem, N, num_levels=4, optimal_trajectories=None,
           local_optimization=True):
    """Generate model inputs using the Method of Morris

    Returns a NumPy matrix containing the model inputs required for Method of
    Morris.  The resulting matrix has :math:`(G+1)*T` rows and :math:`D`
    columns, where :math:`D` is the number of parameters, :math:`G` is the
    number of groups (if no groups are selected, the number of parameters).
    :math:`T` is the number of trajectories :math:`N`,
    or `optimal_trajectories` if selected.
    These model inputs  are intended to be used with
    :func:`SALib.analyze.morris.analyze`.

    Parameters
    ----------
    problem : dict
        The problem definition
    N : int
        The number of trajectories to generate
    num_levels : int, default=4
        The number of grid levels
    optimal_trajectories : int
        The number of optimal trajectories to sample (between 2 and N)
    local_optimization : bool, default=True
        Flag whether to use local optimization according to Ruano et al. (2012)
        Speeds up the process tremendously for bigger N and num_levels.
        If set to ``False`` brute force method is used, unless ``gurobipy`` is
        available

    Returns
    -------
    sample : numpy.ndarray
        Returns a numpy.ndarray containing the model inputs required for Method
        of Morris. The resulting matrix has :math:`(G/D+1)*N/T` rows and
        :math:`D` columns, where :math:`D` is the number of parameters.
    """
    if problem.get('groups'):
        sample = _sample_groups(problem, N, num_levels)
    else:
        sample = _sample_oat(problem, N, num_levels)

    if optimal_trajectories:

        sample = _compute_optimised_trajectories(problem,
                                                 sample,
                                                 N,
                                                 optimal_trajectories,
                                                 local_optimization)

    scale_samples(sample, problem['bounds'])
    return sample

def _sample_oat(problem, N, num_levels=4):
    """Generate trajectories without groups

    Arguments
    ---------
    problem : dict
        The problem definition
    N : int
        The number of samples to generate
    num_levels : int, default=4
        The number of grid levels
    """
    group_membership = np.asmatrix(np.identity(problem['num_vars'],
                                               dtype=int))

    num_params = group_membership.shape[0]
    sample = np.array([generate_trajectory(group_membership,
                                           num_levels)
                       for n in range(N)])
    return sample.reshape((N * (num_params + 1), num_params))

def _sample_groups(problem, N, num_levels=4):
    """Generate trajectories for groups

    Returns an :math:`N(g+1)`-by-:math:`k` array of `N` trajectories,
    where :math:`g` is the number of groups and :math:`k` is the number
    of factors

    Arguments
    ---------
    problem : dict
        The problem definition
    N : int
        The number of trajectories to generate
    num_levels : int, default=4
        The number of grid levels

    Returns
    -------
    numpy.ndarray
    """
    if len(problem['groups']) != problem['num_vars']:
        raise ValueError("Groups do not match to number of variables")

    group_membership, _ = compute_groups_matrix(problem['groups'])

    if group_membership is None:
        raise ValueError("Please define the 'group_membership' matrix")
    if not isinstance(group_membership, np.ndarray):
        raise TypeError("Argument 'group_membership' should be formatted \
                         as a numpy ndarray")

    num_params = group_membership.shape[0]
    num_groups = group_membership.shape[1]
    sample = np.zeros((N * (num_groups + 1), num_params))
    sample = np.array([generate_trajectory(group_membership,
                                           num_levels)
                       for n in range(N)])
    return sample.reshape((N * (num_groups + 1), num_params))

def generate_trajectory(group_membership, num_levels=4):
    """Return a single trajectory

    Return a single trajectory of size :math:`(g+1)`-by-:math:`k`
    where :math:`g` is the number of groups,
    and :math:`k` is the number of factors,
    both implied by the dimensions of `group_membership`

    Arguments
    ---------
    group_membership : np.ndarray
        a k-by-g matrix which notes factor membership of groups
    num_levels : int, default=4
        The number of levels in the grid

    Returns
    -------
    np.ndarray
    """

    delta = compute_delta(num_levels)

    # Infer number of groups `g` and number of params `k` from
    # `group_membership` matrix
    num_params = group_membership.shape[0]
    num_groups = group_membership.shape[1]

    # Matrix B - size (g + 1) * g -  lower triangular matrix
    B = np.tril(np.ones([num_groups + 1, num_groups],
                        dtype=int), -1)

    P_star = generate_p_star(num_groups)

    # Matrix J - a (g+1)-by-num_params matrix of ones
    J = np.ones((num_groups + 1, num_params))

    # Matrix D* - num_params-by-num_params matrix which decribes whether
    # factors move up or down
    D_star = np.diag(rd.choice([-1, 1], num_params))

    x_star = generate_x_star(num_params, num_levels)

    # Matrix B* - size (num_groups + 1) * num_params
    B_star = compute_b_star(J, x_star, delta, B,
                            group_membership, P_star, D_star)

    return B_star

def compute_b_star(J, x_star, delta, B, G, P_star, D_star):
    """
    """
    element_a = J[0, :] * x_star
    element_b = np.matmul(G, P_star).T
    element_c = np.matmul(2.0 * B, element_b)
    element_d = np.matmul((element_c - J), D_star)

    b_star = element_a + (delta / 2.0) * (element_d + J)
    return b_star

def generate_p_star(num_groups):
    """Describe the order in which groups move

    Arguments
    ---------
    num_groups : int

    Returns
    -------
    np.ndarray
        Matrix P* - size (g-by-g)
    """
    p_star = np.eye(num_groups, num_groups)
    rd.shuffle(p_star)
    return p_star

def generate_x_star(num_params, num_levels):
    """Generate an 1-by-num_params array to represent initial position for EE

    This should be a randomly generated array in the p level grid
    :math:`\omega`

    Arguments
    ---------
    num_params : int
        The number of parameters (factors)
    num_levels : int
        The number of levels

    Returns
    -------
    numpy.ndarray
        The initial starting positions of the trajectory

    """
    x_star = np.zeros((1, num_params))
    delta = compute_delta(num_levels)
    bound = 1 - delta
    grid = np.linspace(0, bound, 2)

    x_star[0, :] = rd.choice(grid, num_params)

    return x_star

def compute_delta(num_levels):
    """Computes the delta value from number of levels

    Arguments
    ---------
    num_levels : int
        The number of levels

    Returns
    -------
    float
    """
    return num_levels / (2.0 * (num_levels - 1))

def _compute_optimised_trajectories(problem, input_sample, N, k_choices,
                                    local_optimization=False):
    '''
    Calls the procedure to compute the optimum k_choices of trajectories
    from the input_sample.
    If there are groups, then this procedure allocates the groups to the
    correct call here.

    Arguments
    ---------
    problem : dict
        The problem definition
    input_sample :
    N : int
        The number of samples to generate
    k_choices : int
        The number of optimal trajectories
    local_optimization : bool, default=False
        If true, uses local optimisation heuristic
    '''
    if _has_gurobi is False \
            and local_optimization is False \
            and k_choices > 10:
        msg = "Running optimal trajectories greater than values of 10 \
                will take a long time."
        raise ValueError(msg)

    num_params = problem['num_vars']

    if np.any((input_sample < 0) | (input_sample > 1)):
        raise ValueError("Input sample must be scaled between 0 and 1")

    if _has_gurobi and local_optimization is False:
        # Use global optimization method
        strategy = GlobalOptimisation()
    elif local_optimization:
        # Use local method
        strategy = LocalOptimisation()
    else:
        # Use brute force approach
        strategy = BruteForce()

    if problem.get('groups'):
        num_groups = len(set(problem['groups']))
    else:
        num_groups = None

    context = SampleMorris(strategy)
    output = context.sample(input_sample, N, num_params,
                            k_choices, num_groups)

    return output

def cli_parse(parser):
    parser.add_argument('-l', '--levels', type=int, required=False,
                        default=4, help='Number of grid levels \
                        (Morris only)')
    parser.add_argument('-k', '--k-optimal', type=int, required=False,
                        default=None,
                        help='Number of optimal trajectories \
                        (Morris only)')
    parser.add_argument('-lo', '--local', type=bool, required=True,
                        default=False,
                        help='Use the local optimisation method \
                        (Morris with optimization only)')
    return parser

def cli_action(args):
    rd.seed(args.seed)

    problem = read_param_file(args.paramfile)
    param_values = sample(problem, args.samples, args.levels,
                          args.k_optimal, args.local)

    np.savetxt(args.output, param_values, delimiter=args.delimiter,
               fmt='%.' + str(args.precision) + 'e')

if __name__ == "__main__":
    common_args.run_cli(cli_parse, cli_action)