#!/usr/bin/env python
# -*- coding: utf-8 -*-
# =============================================================================
## @file
#  Module with utilities for reweigting
#
#  @author Vanya BELYAEV Ivan.Belyaev@itep.ru
#  @date   2014-05-10
# =============================================================================
"""Module with utilities for reweighting"""
# =============================================================================
__version__ = "$Revision$"
__author__  = "Vanya BELYAEV Ivan.Belyaev@itep.ru"
__date__    = "2011-06-07"
__all__     = (
    'Weight'          , ## the weighting object
    'makeWeights'     , ## helper function for   weighting iterations
    'WeightingPlot'   , ## helper object to define the weighting rule & target
    'ComparisonPlot'  , ## helper object to keep comparision plot
    'W2Tree'          , ## helper to add the calculated weight to ROOT.TTree
    'W2Data'          , ## helper to add the clacualted weight to ROOT.RooAbsData
    )
# =============================================================================
from   ostap.core.pyrouts     import VE, SE, Ostap
from   ostap.math.base        import iszero
from   ostap.core.core        import split_string
from   ostap.core.ostap_types import string_types, list_types, num_types, sized_types, sequence_types
from   ostap.math.operations  import Mul as MULT  ## needed for proper abstract multiplication
import ostap.io.zipshelve     as     DBASE ## needed to store the weights&histos
from   ostap.trees.funcs      import FuncTree, FuncData ## add weigth to TTree/RooDataSet
from   ostap.utils.utils      import CallThem, is_formula
from   ostap.logger.utils     import pretty_ve
from   ostap.math.reduce      import root_factory
import ostap.histos.histos
import ostap.histos.compare
import ostap.trees.trees
import ostap.fitting.dataset
import ROOT, operator
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger    import getLogger
if '__main__' ==  __name__ : logger = getLogger ( 'ostap.tools.reweight' )
else                       : logger = getLogger ( __name__               )
_new_methods_ = []
# =============================================================================
## @class AttrGetter
#  simple class to bypass operator.attrgetter that
#  has some problem with multiprocessing
class AttrGetter(object):
    """Simple class to bypass operator.attrgetter that
    has some problem with multiprocessing
    """
    def __init__ ( self , *attributes ) :
        self.__attributes = attributes
    def __call__ ( self , obj ) :
        getter = operator.attrgetter( *self.__attributes )
        return getter ( obj )

    @property
    def attributes ( self ) :
        """`attributes': the actual attributes
        """
        return self.__attributes
    # print attributes
    def __str__ ( self ) :
        return ','.join ( self.__attributes )
    __repr__ = __str__
# =============================================================================
## @class Weight
#  helper class for semiautomatic reweighting of data
#  @author Vanya BELYAEV Ivan.Belyaev@itep.ru
#  @date   2014-05-10
class Weight(object) :
    """Helper class for semi-automatic reweighting of data

    It reads various `weigth components' from DBASE and calculates
    the `global' event weight via simple accumulation of weights

    Simplest case: calculate weights for one variable.

    factors = [ ( fun , name ) ]
    where `fun' is a function to get variable from TTree/RooDataSet
    and   `name' is address in data base with reweighting information
    e.g.

    >>> factors = [ ( lambda s : s.pt  , 'pt-data' ) ]

    It assumes, that DBASE has an entry with name 'pt-data', such
    that

    >>> funcs = database['pt-data']

    Here `func' will be function/callable or list of functions/callables
    with weights:

    >>> pt = fun ( entry )
    >>> w  = 1.0
    >>> w *= func ( pt )              ## for the single function/callable
    >>> for f in funcs : w*= f ( pt ) ## for list of functions/callables

    quantity `w' will be an event  weight

    if list of `factors' constains more than one entry,
    to each entry weigth is calculated independently and total weight
    is a product of individual weights.

    For simplistic case one can put some storable function or function object
    >>> db['pt-data'] = math.exp
    >>> db['pt-data'] = [ math.exp , math.cosh ] ## List of object also works

    For realistic case the weights usually stored as (a list of) 1D or 2D histograms
    >>>  db['pt-data'] = [h1,h2,h3,...,hn]
    where h1,...,hn are iteratively obtained histograms with corrections, such as
    the total correction is calculated by the product of them
    """
    ## constructor of
    def __init__ ( self                   ,
                   dbase   = "weights.db" , ## the name of data base with the weights
                   factors = []           ) :

        #
        ## make some statistic
        #
        self.__counter = SE ()
        self.__nzeroes = 0

        self.__vars    = []
        if not factors : return
        self.__dbase   = dbase
        self.__factors = tuple ( factors )

        ## open database

        self.__table = [ ( 'Reweighting' , 'accessor' , '#' , 'merged?' , 'skip') ]
        rows = []

        with DBASE.open ( dbase , 'r' ) as db : ## READONLY

            logger.debug ( 'Weight: Reweigting database: \n%s' % db.table ( prefix = '# ' ) )

            ## loop over the weighting factors and build the function
            for wvar in self.__factors :

                funval  = wvar.accessor  ## accessor to the variable(s)
                funname = wvar.address   ## address  in database
                merge   = wvar.merge     ## merge sequence of callables?
                skip    = wvar.skip      ## skip   some of them?

                row = []

                row.append ( funname )

                if isinstance ( funval , str ) :
                    row.append ( funval )
                    funval = AttrGetter( funval )
                ## elif isinstance ( funval , AttrGetter ) :
                ##    atts = funval.attributes
                ##    if 1 == len ( atts ) : atts = atts[0]
                ##    row.append ( str ( atts ) )
                else :
                    row.append ( str ( funval ) )

                ##
                functions  = db.get ( funname , [] ) ## db[ funname ]
                if not functions :
                    logger.warning ( "Weight: No reweighting is available for `%s', skip it" % funname )
                    continue

                if not isinstance (  functions , sequence_types ) :
                    functions = [ functions ]

                flen = len ( functions )
                if   0 < skip and skip      < flen :
                    logger.info  ("Weight: Use only %d first iterations for `%s'" % ( skip , funname ) )
                    functions = functions[:skip]
                elif 0 > skip and abs(skip) < flen :
                    logger.info  ("Weight: Skip last %d iterations for `%s'" % ( skip , funname ) )
                    functions = functions[:skip]
                elif 0 == skip :
                    pass
                else :
                    logger.error("Weight: Invalid `skip' parameter %s/%d for `%s'" % ( skip , flen , funname ) )
                row.append ( '%d' % flen )

                ## nullify the uncertainties except for the last histogram
                _functions = []
                _first     = True
                for f in reversed ( functions ) :
                    if isinstance ( f , ROOT.TH1 ) and _first :
                        ff = f.clone()
                        for i in ff :
                            v        = float ( ff [ i ] )
                            ff [ i ] = VE ( v , 0 )
                        _functions.append ( ff  )
                        _first = False
                    else :
                        _functions.append ( f   )

                _functions.reverse()
                functions = _functions

                row.append  ( '+' if merge else '-' )
                row.append  ( '%s' % skip           )

                ## merge list of functions into single function
                if merge and 1 < len ( functions)  :

                    ## single_func = functions[0] * functions [1]
                    single_func = MULT ( functions[0] , functions [1] )

                    for fun in functions [2:] :

                        ## multiply it
                        ## single_func *= fun
                        single_func = MULT ( single_func , fun )

                    functions  = [ single_func ]

                self.__vars += [ ( funname , funval , functions , SE() ) ]

                self.__table.append ( row )

        self.__vars = tuple ( self.__vars )

    @property
    def dbase ( self ) :
        """'dbase' : the name of the database for storage of reweigting information
        """
        return self.__dbase

    @property
    def variables ( self ) :
        """'variables' - the structure of reweighting variables/machinery"""
        return self.__vars

    @property
    def stat    ( self ) :
        """'stat' : get the statistic of used weights"""
        return self.__counter
    @property
    def nZeroes ( self ) :
        """'nZeroes': Number of null weights"""
        return self.__nzeroes

    ## calculate the weight for the given "event"
    def __call__ ( self , s ) :
        """Calculate the weigth for the given `event' (==record in TTree/TChain or RooDataSet):
        >>> weight = Weight ( ... )
        >>> tree   = ...
        >>> w = weight ( tree )
        """

        ## initialize the weight
        weight  = VE(1,0)

        ## loop over functions
        for i in self.__vars :

            funval    = i[1] ## accessor
            functions = i[2] ## the functions

            ##  get the weight arguments for the given event
            v       = funval ( s )

            ww = VE(1.0)

            for f in functions :

                if isinstance ( v , tuple ) : w = f ( *v )
                else                        : w = f (  v )

                ww *= w # update the weight factor

            ## keep the statistics
            cnt  = i[3]
            cnt += ww.value()

            ## update the global weight
            weight *= ww

        vw = weight.value()

        self.__counter += vw
        if iszero ( vw ) : self.__nzeroes += 1

        return vw

    ## format information as a table
    def table ( self , title = None , prefix = ' ' ) :
        """Format weighting information as table
        """
        import ostap.logger.table as T
        if title is None : title = "Weight('%s')" % self.__dbase
        return T.table ( self.__table , title = title , prefix = prefix ,
                         alignment = 'llcc' )

    # ===============================================================================
    ## get the dictionary of graphs
    #  for the regular case, it will be graphs that illustrates the convergency
    #  for reweighting iterations
    def graphs ( self ) :
        """Get the dictionary of graphs
        """
        grphs = {}

        import ostap.histos.graphs
        with DBASE.open ( self.dbase , 'r' ) as db : ## READONLY

            for i in self.__vars :

                address   = i[0]
                functions = db.get ( address , () )
                if not functions : continue

                ## grpah with the spread(min/max) weights
                gr1 = ROOT.TGraphAsymmErrors ( len ( functions )  )
                ## graph with the rms weights
                gr2 = ROOT.TGraphErrors      ( len ( functions )  )

                has_points = False
                for n , w  in enumerate ( functions , start = 1 ) :

                    if not hasattr ( w  ,'stat' ) : continue

                    cnt = w.stat()
                    wmean       = cnt.mean  ().value()
                    wrms        = cnt.rms   ()
                    wmn , wmx   = cnt.minmax()
                    gr1 [ n - 1 ] =      n , 0 , 0 , wmean , wmn - wmean , wmx - wmean
                    gr2 [ n - 1 ] = VE ( n , 0.0 ) , VE ( wmean , wrms ** 2)
                    has_points = True

                if has_points :

                    gr1.SetMarkerStyle ( 1  )
                    gr1.SetFillColor   ( ROOT.kOrange )
                    gr1.SetFillStyle   ( 1001 )
                    gr1.SetLineColor   ( ROOT.kOrange )
                    gr1.SetMarkerColor ( ROOT.kOrange )

                    gr2.SetMarkerStyle ( 20 )
                    gr2.SetLineColor   ( 2  )
                    gr2.SetMarkerColor ( 2  )

                    ## combine two graphs togather
                    graph = ROOT.TMultiGraph ()
                    graph.Add      ( gr1 , '3'  )
                    graph.Add      ( gr2 , 'pl' )
                    graph.SetTitle ( '%s;iteration;weights' % address )

                    grphs [ address ] = graph

        return grphs

    def __str__    ( self ) : return self.table()
    def __repr__   ( self ) : return self.table()

    # =========================================================================
    ## reduce the object for serialization (for parallel processing)
    def __reduce__ ( self ) :
        """Reduce the object for serialization (for parallel processing)
        """
        return root_factory , ( type ( self )  ,
                                self.__dbase   ,
                                self.__factors )

    # =========================================================================
    ## @class WeightingVar
    #  Helper class to keep information about single reweighting
    #  - accessor : an accessor function that extracts the variable(s)
    #               from  TTree/TChain/RooDataSet
    #  - address  : the  address in DBASE, where reweighting callable(s) is/are stored
    #  - merge    : merge list of callables from DB into the single callable ?
    #  - skip     : use only certain elements from the list of callables from DBASE
    #
    # Schematic data flow to get the weigth for the given event
    #  - tree/chain/dataset -> accessor -> database(address) -> weight
    class  Var(object) :
        """Helper class to keep information about singe reweighting
        - `accessor'  : an accessor function that extracts the variable(s) from  TTree/TChain/RooDataSet
        - `address'   : the  address in DBASE, where reweigftjnig callable(s) is/are stored
        - `merge'     : merge list of callables from DB into the single callable ?
        - `skip'      : use only certain elements from the list of callables from DBASE

        Schematic data flow to get the weigth for the given event
        - tree/chain/dataset -> accessor -> database(address) -> weight
        """
        ##  constructor
        #   @param sccessor  the accessor function:  tree -> variable(s)
        #   @param address   the address of   reweighintg object in DBASE
        #   @param merge     merge sequence of reweigthing objects ?
        #   @param skip      skip some reweigting objects ?
        def __init__ ( self ,
                       accessor         ,   ## accessor function:  tree -> variable(s)
                       address          ,   ## the address of   reweighting object in DBASE
                       merge     = True ,   ## merge sequence of reweighting objects ?
                       skip      = None ) : ## skip some reweigting objects ?
            """Keep information about single reweighting
            - `accessor' : an accessor function that extracts the variable(s) from  TTree/TChain/RooDataSet
            - `address'  : the  address in DBASE, where reweigftjnig callable(s) is/are stored
            - `merge'    : merge list of callables from DB into the single callable ?
            - `skip'     : use only certain elements from the list of callables from DBASE

            Schematic data flow to get the weigth for the given event
            - tree/chain/dataset -> accessor -> database(address) -> weight
            """

            if isinstance ( accessor , string_types ) :
                accessor = split_string ( accessor , strip = True , respect_groups = True )
                if 1 == len ( accessor ) : accessor = accessor [ 0 ]

            if   accessor and isinstance ( accessor , string_types ) and is_formula ( accessor ) :
                ## use exporession here (suitable both for for TTree/RooAbsData)
                accessor = Ostap.Functions.Expression ( accessor )

            elif accessor and isinstance ( accessor , string_types ) :
                ## accessor = operator.attrgetter (  accessor )
                accessor = AttrGetter (  accessor )

            elif accessor and isinstance ( accessor , sequence_types   ) and \
                     all ( callable ( i ) for i in accessor ) :
                ## merge into single callable
                accessor = CallThem ( accessor )

            elif accessor and isinstance ( accessor , sequence_types   )       and \
                     all ( isinstance ( i , string_types ) for i in accessor ) and \
                     any ( is_formula ( i ) for i in accessor  ) :

                ## convert to universal function for TTree/RooAbsData
                accessor = tuple ( Ostap.Functions.Expression ( i ) for i in accessor )
                ## merge into single callable
                accessor = CallThem ( accessor )

            elif accessor and isinstance ( accessor , sequence_types   )       and \
                     all ( isinstance ( i , string_types ) for i in accessor ) :

                ## accessor = operator.attrgetter ( *accessor )
                accessor = AttrGetter ( *accessor )

            assert accessor and callable ( accessor ) , \
                   "Invalid type of `accessor' %s/%s" % ( accessor , type( accessor ) )

            self.__accessor = accessor ,
            self.__address  = str ( address )
            self.__merge    = True if merge else False
            self.__skip     = skip if skip  else 0

        @property
        def accessor ( self ) :
            """`accessor' - the accessor function to get the variable from TTree/TChain/RooDataSet
            e.g.  to get the variable `x' :
            >>> xvar   = lambda s : s.x
            get to     variables for 2D-reweighting:
            >>> xyvars = lambda s : s.x,s.y
            """
            return self.__accessor[0]

        @property
        def address  ( self ) :
            """`address'  - the address in DB with the reweighting information
            for the  given varibale(s).
            A callable (or list of  callables) is expected in database
            These callables accept their result of `accessor' as their argument
            """
            return self.__address

        @property
        def merge ( self ) :
            """`merge' - merge list  of callables into the single  callable?"""
            return self.__merge

        @property
        def skip ( self ) :
            """`skip'   - use only certain elements from the list of callables from DBASE
            0 < skip ?  - use only the first  `skip' elements
            0 > skip ?  - skip last  `abs(skip)' elements
            """
            return self.__skip

# =============================================================================
## default function to make MC projection
#  @param daatset dataset MC  dataset (typically TTree)
#  @param histo   histogram template
#  @param what    histogram varibales
#  @param how     histogram template
def mc_data_projector  ( dataset , histo , what , how ) :
    """Default function to make MC projection
    """

    dataset.project ( histo , what , how )
    return histo

# =============================================================================
## @class WeightingPlot
#  helper class to manage/keep `weighting-plot'
class WeightingPlot(object) :
    """Helper class to manage/keep `weighting-plot'

    - `what' : the variable/expression to be used for `weighting-plot'
    Used  as the second argument of `dataset.project' method to produce
    `weighted-plot':

    >>> dataset.project ( mchisto , WHAT , how , ... )

    - `how' : the `weight' expression to be used for `weighed-plots'
    Used  as the third argument of `dataset.project' method to produce
    `weighted-plot':

    >>> dataset.project ( mchisto , what , HOW , ... )

    Typically it refers to `weight' variable in the dataset

    - `address' : the addres in `weighting-database'
    where to store the obtained weights

    - `data' : the `data' object, or  the `target' for the reweighting procedure
    Typically it is a histogram. But it could be any kind of callable

    - `mchisto' : template/shape for the mc-histogram, to be used for reweighting.
    It is used as the  first argument of `dataset.project' method

    >>> dataset.project ( MCHISTO , what , how , ... )
    """
    def __init__ ( self              ,
                   what              ,
                   how               ,
                   address           ,
                   data              ,
                   mc_histo  = None  ,
                   w         = 1.0   ,
                   projector = None  ,
                   ignore    = False ) :
        """Helper class to manage/keep `weighting-plot'

        - `what' : the variable/expression to be used for `weighting-plot'
        Used  as the second argument of `dataset.project' method to produce
        `weighted-plot':

        >>> dataset.project ( mchisto , WHAT , how , ... )

        - `how' : the `weight' expression to be used for `weighed-plots'
        Used  as the third argument of `dataset.project' method to produce
        `weighted-plot':

        >>> dataset.project ( mchisto , what , HOW , ... )

        Typically it refers to `weight' variable in datasete

        - `address' : the addres in `weighting-database'
        where to store the obtained weights

        - `data' : the `data' object, or  the `target' for the reweighting procedure
        Typically it is a histotgram. But it could be any kind of callable

        - `mc_histo' : template/shape for the mc-histogram, to be used for reweighting.
        It is used as the  first argument of `dataset.project' method

        >>> dataset.project ( MCHISTO , what , how , ... )
        """

        assert data and callable ( data ) , "ReweightingPlot: `data' object must be callable!"

        self.__what      = str(what)     if isinstance ( what , str ) else what
        self.__how       = str(how )     if isinstance ( how  , str ) else how
        self.__address   = str(address)

        ## ATTENTION! make a ty to convert data to density!!
        if isinstance ( data , ROOT.TH1 ) :
            if not data.is_density() :
                logger.info ( "WeightingPlot: Convert 'data' histogram into 'density' for '%s'" % self.what )
            self.__data = data.density()
        else :
            self.__data = data

        self.__mc       = mc_histo      if mc_histo else data.clone()

        assert isinstance ( self.__mc , ROOT.TH1 ), \
               "WPlot: invalid type of `mchisto' %s/%s"  % ( self.__mc , type ( self.__mc ) )
        self.__w         = w

        if not projector : projector = mc_data_projector

        self.__projector = projector

        assert self.projector and callable ( self.projector ) ,"`Projector' must be callable!"
        self.__ignore    = True if ignore else False

    @property
    def  what ( self ) :
        """'what' : the variable/expression to be used for `weighting-plot'
        Used  as the second argument of `dataset.project' method to produce
        `weighted-plot':
        >>> dataset.project ( mchisto , WHAT , how , ... )
        """
        return self.__what
    @property
    def  how  ( self ) :
        """'how' : the `weight' expression to be used for `weighed-plots'
        Used  as the third argument of `dataset.project' method to produce
        `weighted-plot':
        >>> dataset.project ( mchisto , what , HOW , ... )
        Typically it refers to `weight' variable in datasete
        """
        return self.__how
    @property
    def  address ( self ) :
        """'address' : the address in `weighting-database'
        where to store the obtained weights
        """
        return self.__address
    @property
    def  data ( self ) :
        """'data' : the `data' object, or  the `target' for the reweighting procedure
        Typically it is a histogram. But it could be any kind of callable/function
        """
        return self.__data
    @property
    def  mc_histo ( self ) :
        """'mc_histo' : template/shape for the mc-histogram, to be used for reweighting.
        It is used as the  first argument of `dataset.project' method
        >>> dataset.project ( MCHISTO , what , how , ... )
        """
        return self.__mc
    @property
    def projector ( self ) :
        """'projector' :  callable function to build MC distribution:
        hmc = projector ( dataset , hmc )
        """
        return self.__projector

    @property
    def w  ( self )   :
        """'w'  - relative weight (relative importance is this variable)"""
        return self.__w
    @property
    def ignore ( self ) :
        """'ignore' : do not use variable in reweights, but use for comparsion"""
        return self.__ignore

# =============================================================================
from collections import namedtuple
ComparisonPlot = namedtuple ( 'ComparisonPlot' , ( 'what' , 'data' , 'mc' , 'weight' ) )

_store = set()
# =============================================================================
## draw comparison plots
def _cmp_draw_ ( self ) :
    """Draw comparison plots
    """
    hd   = self.data
    hmc  = self.mc
    hw   = self.weight

    if   isinstance ( hw , ROOT.TH3 ) and 3 == hw.dim () :
        hw.draw  ( copy = True )

        ## minx, miny, minz = hw.minimum_bin()
        ## maxx, maxy, maxz = hw.maximum_bin()

        ## ax = hw.GetXaxis()
        ## ay = hw.GetYaxis()
        ## az = hw.GetZaxis()

        ## minv = hw [ minx, miny, minz ]
        ## maxv = hw [ maxx, maxy, maxz ]

        ## logger.info ( "ComparisonPlot('%s'): min-weight %+.3f at (%.5g,%.5g,%.5g)" % ( self.what ,
        ##                                                                                minv      ,
        ##                                                                                ax.GetBinCenter ( minx ) ,
        ##                                                                                ay.GetBinCenter ( miny ) ,
        ##                                                                                az.GetBinCenter ( minz ) ) )
        ## logger.info ( "ComparisonPlot('%s'): max-weight %+.3f at (%.5g,%.5g,%.5g)" % ( self.what ,
        ##                                                                                maxv      ,
        ##                                                                                ax.GetBinCenter ( maxx ) ,
        ##                                                                                ay.GetBinCenter ( maxy ) ,
        ##                                                                                az.GetBinCenter ( maxz ) ) )

    elif isinstance ( hw , ROOT.TH2 ) and 2 == hw.dim () :

        zmin,zmax = hw.zminmax()

        if   0.95 <= zmin and zmax <= 1.05 :
            hw.SetMinimum  ( 0.95 )
            hw.SetMaximum  ( 1.05 )
            hw.SetContour  ( 50   )
        elif 0.90 <= zmin and zmax <= 1.10 :
            hw.SetMinimum  ( 0.90 )
            hw.SetMaximum  ( 1.10 )
            hw.SetContour  ( 40   )
        elif 0.80 <= zmin and zmax <= 1.20 :
            hw.SetMinimum  ( 0.80 )
            hw.SetMaximum  ( 1.20 )
            hw.SetContour  ( 40   )
        elif 0.75 <= zmin and zmax <= 1.25 :
            hw.SetMinimum  ( 0.75 )
            hw.SetMaximum  ( 1.25 )
            hw.SetContour  ( 50   )
        elif 0.50 <= zmin and zmax <= 1.50 :
            hw.SetMinimum  ( 0.50 )
            hw.SetMaximum  ( 1.50 )
            hw.SetContour  ( 50   )
        else :
            hw.SetMinimum  ( 0.50 )
            hw.SetMaximum  ( 2.00 )
            hw.SetContour  ( 50   )

        nx = hw.binsx()
        ny = hw.binsy()
        if   10 <  nx and 10 < ny : hw.draw  ( 'cont4z' ,                copy = True )
        elif 10 <  nx or  10 < ny : hw.draw  ( 'cont4z' ,                copy = True )
        else                      : hw.texte ( 'colz'   , fmt = '5.3f' , copy = True )

        minx, miny = hw.minimum_bin()
        maxx, maxy = hw.maximum_bin()

        ax = hw.GetXaxis()
        ay = hw.GetYaxis()

        pmin = ROOT.TMarker ( ax.GetBinCenter ( minx ) , ay.GetBinCenter ( miny ) , 43 )
        pmax = ROOT.TMarker ( ax.GetBinCenter ( maxx ) , ay.GetBinCenter ( maxy ) , 20 )
        pmin.SetMarkerColor ( 2 )
        pmax.SetMarkerColor ( 2 )
        pmin.SetMarkerSize  ( 2 )
        pmax.SetMarkerSize  ( 2 )
        pmin.DrawClone()
        pmax.DrawClone()

        ## minv = hw [ minx, miny ]
        ## maxv = hw [ maxx, maxy ]

        ## logger.info ( "ComparisonPlot('%s'): min-weight %+.3f at (%.5g,%.5g)" % ( self.what ,
        ##                                                                           minv      ,
        ##                                                                           ax.GetBinCenter ( minx ) ,
        ##                                                                           ay.GetBinCenter ( miny ) ) )
        ## logger.info ( "ComparisonPlot('%s'): max-weight %+.3f at (%.5g,%.5g)" % ( self.what ,
        ##                                                                           maxv      ,
        ##                                                                           ax.GetBinCenter ( maxx ) ,
        ##                                                                           ay.GetBinCenter ( maxy ) ) )

    elif isinstance ( hw , ROOT.TH1 ) and 1 == hw.dim () :
        if hasattr ( hw     , 'red'   ) : hw .red   ()
        hw.SetMinimum  ( 0 )
        hw.SetMaximum  ( max ( 1.3 , 1.2 * hw.ymax () ) )
        hw.draw ( copy = True )
        hw.level ( 1.0 , linestyle = 1 , linecolor = ROOT.kRed )

    if isinstance ( hd , ROOT.TH1 ) and isinstance ( hmc , ROOT.TH1 ) and 1 == hd.dim() :

        hd  = hd .clone()
        hmc = hmc.clone()

        if hasattr ( hd     , 'green'  ) : hd .green   ()
        if hasattr ( hmc    , 'blue'   ) : hmc.blue   ()
        rmax   = 1.2 * max ( hd .GetMaximum () , hmc.GetMaximum() )
        hd .SetMinimum ( 0 )
        hmc.SetMinimum ( 0 )
        scale  = ROOT.gPad.GetUymax() / rmax
        hd  .Scale ( scale  )
        hmc .Scale ( scale  )
        hd  .draw  ( 'same' , copy = True )
        hmc .draw  ( 'same' , copy = True )

    elif isinstance ( hd  , ROOT.TH1 ) and 1 == hd.dim() :

        hd = hd.clone()
        if hasattr ( hd     , 'green'  ) : hd .green   ()
        rmax   = 1.2 *       hd .GetMaximum ()
        hd .SetMinimum ( 0 )
        scale  = ROOT.gPad.GetUymax() / rmax
        hd  .Scale ( scale  )
        hd  .draw  ( 'same' , copy = True )

    elif isinstance ( hmc , ROOT.TH1 ) and 1 == hmc.dim() :

        hmc = hmc.clone()

        if hasattr ( hmc    , 'blue'   ) : hmc.blue   ()
        rmax   = 1.2 *       hmc.GetMaximum ()
        hmc.SetMinimum ( 0 )
        scale  = ROOT.gPad.GetUymax() / rmax
        hmc .Scale ( scale  )
        hmc .draw  ( 'same' , copy = True )

ComparisonPlot. draw = _cmp_draw_

# =============================================================================
## The main  function: perform one re-weighting iteration
#  and reweight "MC"-data set to looks as "data"(reference) dataset
#  @code
#  results = makeWeights (
#   dataset              , ## data source to be  reweighted (DataSet, TTree, abstract source)
#   plots                , ## reweighting plots
#   database             , ## datadabse to store/update reweigting results
#   delta                , ## stopping criteria for "mean"    weight variation
#   minmax               , ## stopping criteria for "min/max" weight variation
#   power                , ## effective power to apply to the weigths
#   debug        = True  , ## store debuig information in database
#   make_plots   = True  , ## produce useful comparison plots
#   force_update = False , ## force DB update even for "good" results
#   tag          = 'RW'  ) ## tag for better printout
#  @endcode
#  If make_plots = False function returns the tuple of active reweitings:
#  @code
#  active        = makeWeights ( ... , make_plots = False , ... ) \
#  @endcode
#  otherwise it also returns the tuple of comparison plots
#  @code
#  active, cmp_plots = makeWeights ( ... , make_plots = True  , ... )
#  for item in  cmp_plots :
#     what    = item.what
#     hdata   = item.data
#     hmc     = item.mc
#     hweight = item.weight
#  @endcode
#  If no more rewighting iteratios required, active is an empty tuple
#  @author Vanya BELYAEV Ivan.Belyaev@itep.ru
#  @date   2014-05-10
def makeWeights  ( dataset                    ,
                   plots        = []            ,
                   database     = "weights.db"  ,
                   compare      = None          , ## comparison function
                   delta        = 0.01          , ## delta for `mean'  weight variation
                   minmax       = 0.03          , ## delta for `minmax' weight variation
                   power        = None          , ## auto-determination
                   debug        = True          , ## save intermediate information in DB
                   make_plots   = True          , ## make comparison plots (and draw them)
                   wtruncate    = ( 0.5 , 1.5 ) , ## truncate too small/large weights
                   force_update = False         , ## force DB update  ven for "good" results
                   tag          = "Reweighting" ) :
    """The main  function: perform one re-weighting iteration
    and reweight `MC'-data set to looks as `data'(reference) dataset
    >>> results = makeWeights (
    ... dataset           , ## data source to be  reweighted (DataSet, TTree, abstract source)
    ... plots             , ## reweighting plots
    ... database          , ## datadabse to store/update reweigting results
    ... delta             , ## stopping criteria for `mean`    weight variation
    ... minmax            , ## stopping criteria for `min/max` weight variation
    ... power             , ## effective power to apply to the weigths
    ... debug      = True , ## store debuig information in database
    ... make_plots = True , ## produce useful comparison plots
    ... tag        = 'RW' ) ## tag for better printout

    If `make_plots = False`,  it returns the tuple of active reweitings:
    >>> active        = makeWeights ( ... , make_plots = False , ... )

    Otherwise it also returns list of comparison plots
    >>> active, cmp_plots = makeWeights ( ... , make_plots = True  , ... )
    >>> for item in  cmp_plots :
    ...    what    = item.what
    ...    hdata   = item.data
    ...    hmc     = item.mc
    ...    hweight = item.weight

    If no more rewighting iteratios required, active is an empty tuple
    """

    assert 0 < delta  , "makeWeights(%s): Invalid value for `delta'  %s" % ( tag , delta  )
    assert 0 < minmax , "makeWeights(%s): Invalid value for `minmax' %s" % ( tag , minmax )

    if wtruncate and isinstance ( wtruncate , sized_types ) and 2 == len ( wtruncate ) \
       and isinstance ( wtruncate [ 0 ] , num_types ) \
       and isinstance ( wtruncate [ 1 ] , num_types ) \
       and wtruncate [0] < 1.0 and 1.0 < wtruncate [1] : pass
    else :
        logger.warning  ( "Invalid 'wtruncate' parameter '%s', use (0.5,2.0)" % ( str ( wtruncate ) ) )
        wtruncate = 0.5 , 2.0

    from ostap.logger.colorized   import allright , attention , infostr
    from ostap.utils.basic        import isatty

    nplots  = len ( plots )

    for_update   = 0
    ## list of plots to compare
    cmp_plots    = []
    ## reweighting summary table
    header       = ( 'Reweighting' , 'wmin/wmax' , 'OK?' , 'wrms[%]' , 'OK?' , 'chi2/ndf' , 'ww' , 'exp' )
    rows         = {}
    save_to_db   = []
    ## number of active plots for reweighting
    for wplot in plots  :

        what      = wplot.what       ## variable/function to plot/compare
        how       = wplot.how        ## weight and/or additional cuts
        address   = wplot.address    ## address in database
        hdata0    = wplot.data       ## original "DATA" object
        hmc0      = wplot.mc_histo   ## original "MC"   histogram
        ww        = wplot.w          ## relative weight
        projector = wplot.projector  ## projector for MC data
        ignore    = wplot.ignore     ## ignore for weight building?
        #
        # normalize the data
        #
        hdata = hdata0
        if isinstance ( hdata , ROOT.TH1 ) and not hdata.is_density() :
            hdata = hdata.density ()
            logger.warning ( "'Data' histogram converted to 'density' for '%s'" %  what )

        # =====================================================================
        ## make a plot on (MC) data with the weight
        # =====================================================================
        hmc0 = projector ( dataset , hmc0 , what , how )

        st   = hmc0.stat()
        mnmx = st.minmax()
        if  iszero ( mnmx [0] ) :
            logger.warning ( "%s: statistic goes to zero %s/`%s'" % ( tag , st , address ) )
        elif mnmx [0] <= 0      :
            logger.warning ( "%s: statistic is negative  %s/`%s'" % ( tag , st , address ) )

        # =====================================================================
        ## normalize MC
        # =====================================================================
        if isinstance ( hmc0 , ROOT.TH1 ) or hasattr ( hmc , 'density' ) :
            hmc = hmc0.density()
        else : hmc = hmc0

        # =====================================================================
        ## calculate  the reweighting factor : a bit conservative (?)
        #  this is the only important line
        # =====================================================================

        #  try to exploit finer binning if/when possible
        hboth = isinstance ( hmc   , ROOT.TH1 ) and isinstance ( hdata , ROOT.TH1 )

        if   hboth and 1 == hmc.dim () and 1 == hdata.dim () and \
               len ( hmc ) >= len( hdata ) :
            w = ( 1.0 / hmc ) * hdata                                 ## NB!
        elif hboth and 2 == hmc.dim () and 2 == hdata.dim () and \
                 ( hmc.binsx() >= hdata.binsx() ) and \
                 ( hmc.binsy() >= hdata.binsy() ) and \
                 ( hmc.binsx() >  hdata.binsx()   or  \
                   hmc.binsy() >  hdata.binsy() ) :
            w = ( 1.0 / hmc ) * hdata                                 ## NB!
        elif hboth and 3 == hmc.dim () and 3 == hdata.dim () and \
                 ( hmc.binsx() >= hdata.binsx() ) and \
                 ( hmc.binsy() >= hdata.binsy() ) and \
                 ( hmc.binsz() >= hdata.binsz() ) and \
                 ( hmc.binsx() >  hdata.binsx()   or  \
                   hmc.binsy() >  hdata.binsy()   or  \
                   hmc.binsz() >  hdata.binsz() ) :
            w = ( 1.0 / hmc ) * hdata                                 ## NB!
        else                            :
            w = hdata / hmc                                           ## NB!

        # =====================================================================
        ## scale & get the statistics of weights
        w   /= w.stat().mean().value()
        cnt  = w.stat()
        #
        mnw , mxw = cnt.minmax()
        wvar  = cnt.rms()/cnt.mean()
        good1 = wvar.value()      <= delta
        good2 = abs ( mxw - mnw ) <= minmax
        good  = good1 and good2  ## small variance?
        #

        c2ndf  = 0
        for i in w : c2ndf += w[i].chi2 ( 1.0 )
        c2ndf /= ( len ( w ) - 1 )

        ## build  the row in the summary table
        row = [ address  , '%-5.3f/%5.3f' % ( cnt.minmax()[0]    , cnt.minmax()[1] ) ]

        if   ignore : row.append ( ''                )
        elif good2  : row.append ( allright  ( '+' ) )
        else        : row.append ( attention ( '-' ) )

        row.append ( (wvar * 100).toString('%6.2f+-%-6.2f') )

        if   ignore : row.append ( ''                )
        elif good1  : row.append ( allright  ( '+' ) )
        else        : row.append ( attention ( '-' ) )

        row.append (  '%6.2f' % c2ndf  )

        ## apply weight truncation:
        if not ignore :
            wmin , wmax = wtruncate
            truncated = False
            for i in w  :
                we = w [ i ]
                if   we.value() < wmin :
                    w [ i ] = VE ( wmin , we.cov2 () )
                    truncated    = True
                elif we.value() > wmax :
                    w [ i ] = VE ( wmax , we.cov2 () )
                    truncated    = True

            if truncated :
                logger.info ( "%s: weights for '%s' are truncated to be %.3f<w<%.3f" % ( tag , what , wmin , wmax ) )

        ## make plots at the start of  each iteration?
        if make_plots :
            item = ComparisonPlot ( what , hdata , hmc , w )
            cmp_plots.append ( item )

        if ignore : row.append ( '' )
        else      : row.append ( '%4.3f' % ww if 1 != ww  else '' )

        rows [ address ] = tuple ( row )

        #
        ## make decision based on the variance of weights
        #
        mnw , mxw = cnt.minmax()
        if not ignore  :
            ## update DB for "not-good" or "forced" entries
            if force_update or not good :
                if not good : for_update += 1
                save_to_db.append ( ( address , ww , hdata0 , hmc0 , hdata , hmc , w ) )

        # =====================================================================
        ## make a comparison (if needed)
        # =====================================================================
        if compare : compare ( hdata0 , hmc0 , address )

    active  = tuple ( [ p [ 0 ] for p in save_to_db ] )
    nactive = len ( active )

    if   power is None :
        power   = lambda n : 0.5 * ( 1.0 / max ( n , 1 ) + 1 )
        ## average between 100% correlated and 100% uncorrelated
        eff_exp = 0.5 * ( 1.0 / max ( nactive , 1 ) + 1.0 )
    elif power and callable ( power ) :
        eff_exp = power ( nactive )
    elif isinstance ( power , num_types ) and 0 < power <= 1.5 :
        eff_exp = 1.0 * power
    elif 1 == nactive and 1 < len ( plots ) :
        eff_exp = 0.95
    elif 1 == nactive  :
        eff_exp = 1.00
    else               :
        eff_exp = 1.10 / max ( nactive , 1 )

    while database and save_to_db :

        entry = save_to_db.pop()

        address , ww , hd0 , hm0 , hd , hm , weight = entry

        cnt = weight.stat()
        mnw, mxw = cnt.minmax()

        if 1 < nactive and 1 != ww :
            eff_exp *= ww
            logger.info  ("%s: apply `effective exponent' of %.3f for `%s'" % ( tag , eff_exp  , address ) )

        if 1 != eff_exp and 0 < eff_exp :
            weight = weight ** eff_exp

            row = list ( rows [ address ] )
            row.append ( '%4.3f' % eff_exp )
            rows [ address ] = tuple ( row )

        with DBASE.open ( database ) as db :

            db [ address ] = db.get( address , [] ) + [ weight ]

            ## save more information to database
            if debug :
                addr        = address + ':REWEIGHTING'
                db [ addr ] = db.get ( addr , [] ) + [ entry [ 2 : ]  ]
                ## if make_plots :
                ##     found = False
                ##     for c in cmp_plots :
                ##         if c.what == address :
                ##             found = c
                ##             break
                ##     if found :

        del hd0, hm0 , hd , hm , weight , entry

    table = [  header ]
    for   row in   rows  : table.append ( rows[row] )

    import ostap.logger.table as Table
    logger.info ( '%s: %d active reweightings\n%s' % ( tag , nactive , Table.table ( table , title = tag , prefix = '# ' , alignment = 'lccccccc' ) ) )

    cmp_plots = tuple ( cmp_plots )

    if make_plots and cmp_plots and not ROOT.gROOT.IsBatch() :
        from ostap.plotting.canvas import use_canvas
        for item in cmp_plots :
            with use_canvas ( '%s/%s' % ( tag , item.what ) ) : item.draw()

    ## return ( active , cmp_plots ) if make_plots else active
    return ( for_update , cmp_plots ) if make_plots else for_update

# =============================================================================
## @class W2Tree
#  Helper class to add the weight into ROOT.TTree
#  @code
#  w    = Weight ( ... )      ## the weighting object
#  tree = ...                 ## The tree
#  wf   = W2Tree ( w , tree ) ## create the weighting function
#  tree.add_new_branch ( 'weight' , wf )
#  @endcode
class W2Tree(FuncTree) :
    """Helper class to add the weight into ROOT.TTree
    >>> w    = Weight ( ... )      ## the weighting object
    >>> tree = ...                 ## The tree
    >>> wf   = W2Tree ( w , tree ) ## create the weighting function
    >>> tree.add_new_branch ( 'weight' , wf )
    """
    def __init__ ( self , weight , tree = None ) :

        assert isinstance ( weight , Weight    )                 , 'Wrong type of weight!'
        assert tree is None or isinstance ( tree  , ROOT.TTree ) , 'Wrong type of tree!'

        FuncTree.__init__ ( self , tree )
        self.__weight = weight

    ## evaluate the weighter for the given TTree entry
    def evaluate ( self ) :
        """Evaluate the weighter for the given TTree entry"""
        t = self.tree ()
        w = self.__weight ( t )
        return w

    @property
    def weight ( self ) :
        """`weight' : get the weighter object"""
        return  self.__weight

# =============================================================================
## @class W2Data
#  Helper class to add the weight into RooDataSet
#  @code
#  w    = Weight ( ... )      ## the weighting object
#  ds   = ...                 ## dataset
#  wf   = W2Data ( w , tree ) ## create the weighting function
#  ds.add_new_var ( 'weight' , wf )
#  @endcode
class W2Data(FuncData) :
    """Helper class to add the weight into RooDataSet
    >>> w    = Weight ( ... )      ## the weighting object
    >>> ds   = ...                 ## dataset
    >>> wf   = W2Data ( w , tree ) ## create the weighting function
    >>> ds.add_new_var ( 'weight' , wf )
    """

    def __init__ ( self , weight , data = None ) :

        assert isinstance ( weight , Weight    )                      , 'Wrong type of weight!'
        assert data is None or isinstance ( data  , ROOT.RooAbsData ) , 'Wrong type of data!'

        FuncData.__init__ ( self , data )
        self.__weight = weight

    ## evaluate the weighter for the given RooAbsData entry
    def evaluate ( self ) :
        """Evaluate the weighter for the given RooAbsData entry"""

        d = self.data ()
        w = self.__weight ( d )
        return w

    @property
    def weight ( self ) :
        """`weight' : get the weighter object"""
        return  self.__weight

# =============================================================================
## Add specific re-weighting information into ROOT.TTree
#  @see ostap.tools.reweight
#  @see ostap.tools.reweight.Weight
#  @see ostap.tools.reweight.W2Tree
#  @code
#  w    = Weight ( ... ) ## weighting object ostap.tools.reweight.Weight
#  tree = ...
#  tree.add_reweighting ( w )
#  @endcode
def tree_add_reweighting ( tree                 ,
                           weighter             ,
                           name      = 'weight' ,
                           verbose   = True     ,
                           report    = True     ) :
    """Add specific re-weighting information into ROOT.TTree

    >>> w    = Weight ( ... ) ## weighting object ostap.tools.reweight.Weight
    >>> data = ...
    >>> data.add_reweighting ( w )
    - see ostap.tools.reweight
    - see ostap.tools.reweight.Weight
    - see ostap.tools.reweight.W2Tree
    """

    assert isinstance ( weighter , Weight ), "Invalid type of `weighter'!"

    ## create the weighting function
    wfun = W2Tree ( weighter )

    return tree.add_new_branch (  name , wfun , verbose = verbose  , report = report )

ROOT.TTree.add_reweighting = tree_add_reweighting

# =============================================================================
## Add specific re-weighting information into dataset
#  @see ostap.tools.reweight
#  @see ostap.tools.reweight.Weight
#  @see ostap.tools.reweight.W2Data
#  @code
#  w    = Weight ( ... ) ## weighting object ostap.tools.reweight.Weight
#  data = ...
#  data.add_reweighting ( w )
#  @endcode
def data_add_reweighting ( data , weighter , name = 'weight' , progress = False ) :
    """Add specific re-weighting information into dataset

    >>> w    = Weight ( ... ) ## weighting object ostap.tools.reweight.Weight
    >>> data = ...
    >>> data.add_reweighting ( w )
    - see ostap.tools.reweight
    - see ostap.tools.reweight.Weight
    - see ostap.tools.reweight.W2Data
    """

    assert isinstance ( weighter , Weight ), "Invalid type of `weighter'!"

    ## create the weigthting function
    wfun = W2Data ( weighter  )

    if progress :
        from ostap.utils.progress_conf import progress_conf
        return data.add_new_var ( name , wfun , progress_conf() )

    return data.add_new_var ( name , wfun )

ROOT.RooDataSet.add_reweighting = data_add_reweighting

_new_methods_ += [
    ROOT.RooDataSet .add_reweighting ,
    ROOT.TTree      .add_reweighting ,
    ]

# =============================================================================
if '__main__' == __name__ :

    from ostap.utils.docme import docme
    docme ( __name__ , logger = logger )

# =============================================================================
##                                                                      The END
# =============================================================================


import model_zoo
import torch

# ======================================================================
# brief description of this experiment
# ======================================================================
# ...

# ======================================================================
# Model settings
# ======================================================================
model_handle = model_zoo.UNet
use_adaptive_batch_norm = True # If True set train_with_Bern to False

train_with_bern = False
defrozen_conv_blocks = False

# If cut_z is True, we cut the first and last 3 slices of the z axis
cut_z = False
cut_z_saved = 3
run_number = 1
da_ratio = 0.0
nchannels = 4 # [intensity, vx, vy, vz]
#note = f'_bern_{train_with_bern}_adaptive_{use_adaptive_batch_norm}'
note = f'_full_run'
extra_info = "_only_w_labels_adaptive_batch_norm_e80_lr_1e-3_AdamW_val_40"

use_saved_model = True
with_validation = True
train_file_name = 'only_w_labels_bern_images_and_labels_from_101_to_127' # bern_images_and_labels_from_101_to_116, only_w_labels_bern_images_and_labels_from_101_to_116, size_32_bern_images_and_labels_from_101_to_104
# If with_validation is True, we use the following file for validation
# else it's loaded but not used
val_file_name = 'size_40_bern_images_and_labels_from_122_to_127'
add_pixels_weight = 100
# ======================================================================
# data settings
# ======================================================================
data_mode = '3D'
image_size = [144, 112, 48] # [x, y, time]
nlabels = 2 # [background, foreground]

# ======================================================================
# training settings
# ======================================================================
#max_steps = 10000
#steps = 1000
epochs = 2
batch_size = 8
learning_rate = 1e-3
optimizer_handle = torch.optim.AdamW
betas = (0.9, 0.98)
loss_type = 'dice'  # crossentropy/dice
summary_writing_frequency = 20 # 4 times in each epoch (if n_tr_images=20, batch_size=8)
train_eval_frequency = 500 # every 4 epochs (if n_tr_images=20, batch_size=8)
val_eval_frequency = 500 # every 4 epochs (if n_tr_images=20, batch_size=8)
save_frequency = 200 # every 10 epochs (if n_tr_images=20, batch_size=8)

continue_run = False
augment_data = False
experiment_name_saved_model = 'unet3d_da_' + str(da_ratio) + 'nchannels' + str(nchannels) + '_r' + str(run_number) + '_loss_' + loss_type + '_cut_z_' + str(cut_z_saved) + note
experiment_name = 'unet3d_da_' + str(da_ratio) + 'nchannels' + str(nchannels) + '_r' + str(run_number) + '_loss_' + loss_type + '_cut_z_' + str(cut_z) + note + extra_info


from time import time

import numpy as np
import pandas as pd

# utilities
from sklearn.grid_search import RandomizedSearchCV

# models
from sklearn.linear_model import Perceptron, LogisticRegression, BayesianRidge, SGDClassifier, \
    PassiveAggressiveClassifier
from sklearn.svm import LinearSVC, SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier

# more information about sklearn's incremental models
# http://scikit-learn.org/stable/auto_examples/linear_model/plot_sgd_comparison.html
# http://scikit-learn.org/stable/modules/scaling_strategies.html
# TODO add xgboost, keras
models_linear = [Perceptron, LogisticRegression, BayesianRidge, LinearSVC]
models_online = [Perceptron, MultinomialNB, SGDClassifier, PassiveAggressiveClassifier]
models_nonlinear_cheap = [DecisionTreeClassifier]
models_nonlinear_expensive = [RandomForestClassifier, SVC, GradientBoostingClassifier]

hyperparameters = {
    LogisticRegression: {'C': [1, 10, 50, 100, 1000, 10000]},
  # [{'C': [1, 10, 100, 1000], 'kernel': ['linear']},
  # {'C': [1, 10, 100, 1000], 'gamma': [0.001, 0.0001], 'kernel': ['rbf']},]
    RandomForestClassifier: {"max_depth": [3, None],
              "max_features": [1, 3, 10],
              "min_samples_split": [1, 3, 10],
              "min_samples_leaf": [1, 3, 10],
              "bootstrap": [True, False],
              "criterion": ["gini", "entropy"]},
}

# parameters
random_state = 0
random_search_iterations_max = 10

def requirements_bare_minimum(y_train):
    # by default model has to beat:
    # 1- random guessing (AUC must be higher than 0.5), and
    # 2- majority vote classification, important for imbalanced datasets, and
    # 3- cross-validation variation should not be too high (less than 0.1)

    return {'auc': 0.5,
            'accuracy': pd.Series(y_train).value_counts(normalize=True)[0],
            'accuracy_std': 0.1}

def check_requirements(model, requirements):
    if model_insights['accuracy'] > requirements['accuracy'] and \
                        model_variation < requirements['accuracy_std'] and \
                        model_insights['auc'] > requirements['auc'] and \
                        model_insights['accuracy_test'] > requirements['accuracy']:
        return True
    else:
        return False

def tune(insights, x_train, y_train, x_test, y_test, models='all', requirements=None, maximize=False):
    if requirements is None:
        requirements = requirements_bare_minimum(y_train)

    # do vanilla models satisfy the requirements?
    # assuming decision tree is the most intuitive, then logistic regression and then random forest
    # TODO: extend this to metrics other than accuracy using the confusion matrix
    for model_name in ['dt', 'lr', 'rf']:
        model_insights = insights[model_name]
        model_variation = np.std(model_insights['accuracy_folds'])

        if check_requirements(model_insights, requirements) and not maximize:
            pass
            # TODO: turn this back on
            # return model_name

    # model selection and tuning loop
    models_to_train = []

    if models == 'all':
        models_to_train += models_linear + models_nonlinear_cheap + models_nonlinear_expensive
    elif models == 'linear':
        models_to_train += models_online
    elif models_to_train == 'cheap':
        models_to_train += models_linear + models_nonlinear_cheap

    # TODO: using all of the training data, need to use less data if runtime for insights models is large (how large?)
    for model in models_to_train:
        # TODO: add the looping logic
        if model == LogisticRegression:
            number_configurations = np.prod(np.array([len(_) for _ in hyperparameters[model]]))
            random_search_iterations = np.min([random_search_iterations_max, number_configurations])
            random_search = RandomizedSearchCV(model(n_jobs=-1, random_state=random_state),
                param_distributions=hyperparameters[model], n_iter=random_search_iterations, n_jobs=-1, random_state=0)
            runtime = time()
            random_search.fit(x_train, y_train)
            runtime = time() - runtime

            info = dict()
            info['runtime'] = runtime
            # info['accuracy'] = min(scores)
            # info['accuracy_test'] = accuracy_score(y_test, y_test_predicted)
            # info['accuracy_folds'] = scores
            # info['confusion_matrix'] = confusion_matrix(y_test, y_test_predicted)
            # clf.fit(x_train, y_train)
            # fpr, tpr, _ = roc_curve(y_test, clf_predict_proba(clf, x_test))
            # info['fpr'] = fpr
            # info['tpr'] = tpr
            # info['auc'] = auc(fpr, tpr)

            return random_search

    return None


# coding=utf-8
# !/usr/bin/python
import sys
sys.path.append('..')
from base.spider import Spider
import json
import re
import urllib

class Spider(Spider):  # 元类 默认的元类 type
    def getName(self):
        return "Alist"

    def init(self, extend=""):
        print("============{0}============".format(extend))
        pass

    def isVideoFormat(self, url):
        pass

    def manualVideoCheck(self):
        pass

    def homeContent(self, filter):
        result = {}
        cateManual = {
          "菊花盘": "https://pan.142856.xyz/OneDrive",
		"课程搜索": "https://www.souke365.cn",
		"博哥": "www.top12345.top",
          "🔮嗨翻":"https://pan.hikerfans.com",
		  "帅仔1":"http://43.139.29.179:5244/",
                  "帅仔2":"http://43.139.29.179:5245/",
		  "🦀9T(Adult)":"https://drive.9t.ee",
		  "🐱梓澪の妙妙屋":"https://xn--i0v44m.xyz",
		  "🚆资源小站":"https://pan.142856.xyz",
		  "🌤晴园的宝藏库":"https://alist.52qy.repl.co",
		  "🐭米奇妙妙屋":"https://anime.mqmmw.ga",
		  "💂小兵组网盘影视":"https://6vv.app",
		  "📀小光盘":"https://alist.xiaoguanxiaocheng.life",
		  "🐋一只鱼":"https://alist.youte.ml",
		  "🌊七米蓝":"https://al.chirmyram.com",
		  "🌴非盘":"http://www.feifwp.top",
		  "🥼帅盘":"https://hi.shuaipeng.wang",
		  "🐉神族九帝":"https://alist.shenzjd.com",
		  "☃姬路白雪":"https://pan.jlbx.xyz",
		  "🎧听闻网盘":"https://wangpan.sangxuesheng.com",
		  "💾DISK":"http://124.222.140.243:8080",
		  "🌨云播放":"https://quanzi.laoxianghuijia.cn",
		  "✨星梦":"https://pan.bashroot.top",
		  "🌊小江":"https://dyj.me",
		  "💫触光":"https://pan.ichuguang.com",
		  "🕵好汉吧":"https://8023.haohanba.cn",
		  "🥗AUNEY":"http://121.227.25.116:8008",
		  "🎡资源小站":"https://960303.xyz/",
		  "🐝神器云": "https://quanzi.laoxianghuijia.cn",
		  "🏝fenwe":"http://www.fenwe.tk:5244",
		  "🎢轻弹浅唱":"https://g.xiang.lol"
        }
        classes = []
        for k in cateManual:
            classes.append({
                'type_name': k,
				"type_flag": "1",
                'type_id': cateManual[k]
            })
        result['class'] = classes
        if (filter):
            result['filters'] = self.config['filter']
        return result

    def homeVideoContent(self):
        result = {
            'list': []
        }
        return result

    ver = ''
    baseurl = ''
    def getVersion(self, gtid):
        param = {
            "path": '/'
        }
        if gtid.count('/') == 2:
            gtid = gtid + '/'
        baseurl = re.findall(r"http.*://.*?/", gtid)[0]
        ver = self.fetch(baseurl + 'api/public/settings', param)
        vjo = json.loads(ver.text)['data']
        if type(vjo) is dict:
            ver = 3
        else:
            ver = 2
        self.ver = ver
        self.baseurl = baseurl

    def categoryContent(self, tid, pg, filter, extend):
        result = {}
        if tid.count('/') == 2:
            tid = tid + '/'
        nurl = re.findall(r"http.*://.*?/", tid)[0]
        if self.ver == '' or self.baseurl != nurl:
            self.getVersion(tid)
        ver = self.ver
        baseurl = self.baseurl
        if tid.count('/') == 2:
            tid = tid + '/'
        pat = tid.replace(baseurl,"")
        param = {
            "path": '/' + pat
        }
        if ver == 2:
            rsp = self.postJson(baseurl + 'api/public/path', param)
            jo = json.loads(rsp.text)
            vodList = jo['data']['files']
        elif ver == 3:
            rsp = self.postJson(baseurl + 'api/fs/list', param)
            jo = json.loads(rsp.text)
            vodList = jo['data']['content']

        videos = []
        cid = ''
        purl = ''
        svodList = str(vodList)
        lenvodList = len(vodList)
        substr = str(re.findall(r"\'name\': \'(.*?)\'", svodList))
        foldernum = svodList.count('\'type\': 1')
        filenum = lenvodList - foldernum
        for vod in vodList:
            if ver == 2:
                img = vod['thumbnail']
            elif ver == 3:
                img = vod['thumb']
            if len(img) == 0:
                if vod['type'] == 1:
                    img = "http://img1.3png.com/281e284a670865a71d91515866552b5f172b.png"
            if pat != '':
                aid = pat + '/'
            else:
                aid = pat
            if vod['type'] == 1:
                tag = "folder"
                remark = "文件夹"
                cid = baseurl + aid + vod['name']
            #计算文件大小
            else:
                size = vod['size']
                if size > 1024 * 1024 * 1024 * 1024.0:
                    fs = "TB"
                    sz = round(size / (1024 * 1024 * 1024 * 1024.0), 2)
                elif size > 1024 * 1024 * 1024.0:
                    fs = "GB"
                    sz = round(size / (1024 * 1024 * 1024.0), 2)
                elif size > 1024 * 1024.0:
                    fs = "MB"
                    sz = round(size / (1024 * 1024.0), 2)
                elif size > 1024.0:
                    fs = "KB"
                    sz = round(size / (1024.0), 2)
                else:
                    fs = "KB"
                    sz = round(size / (1024.0), 2)
                tag = "file"
                remark = str(sz) + fs
                cid = baseurl + aid + vod['name']
                # 开始爬视频与字幕
                if filenum < 150:
                    if 'mp4' in vod['name'] or 'mkv' in vod['name'] or 'TS' in vod['name'] or 'flv' in vod['name'] or 'rmvb' in vod['name'] or 'mp3' in vod['name'] or 'flac' in vod['name'] or 'wav' in vod['name'] or 'wma' in vod['name'] or 'dff' in vod['name']:
                        vodurl = vod['name']
                        # 开始爬字幕
                        cid = '###'
                        subname = re.findall(r"(.*)\.", vod['name'])[0]
                        if filenum == 2:
                            if '.ass' in substr:
                                sub = re.findall(r"'(.*).ass", substr)[0]
                                subt = '@@@' + sub + '.ass'
                                if ',' in sub:
                                    sub = re.findall(r"', '(.*).ass", substr)[0]
                                    subt = '@@@' + sub + '.ass'
                            if '.srt' in substr:
                                sub = re.findall(r"'(.*).srt", substr)[0]
                                subt = '@@@' + sub + '.srt'
                                if ',' in sub:
                                    sub = re.findall(r"', '(.*).srt", substr)[0]
                                    subt = '@@@' + sub + '.srt'
                        else:
                            if subname + '.ass' in substr:
                                subt = '@@@' + subname + '.ass'
                            elif subname + '.srt' in substr:
                                subt = '@@@' + subname + '.srt'
                        # 合并链接
                        if 'subt' in locals().keys():
                            purl = purl + '{0}{1}#'.format(vodurl, subt)
                        else:
                            purl = purl + '{0}#'.format(vodurl)
                else:
                    subname = re.findall(r"(.*)\.", vod['name'])[0]
                    if subname + '.ass' in substr:
                        subt = '@@@' + subname + '.ass'
                        cid = cid + subt
                    elif subname + '.srt' in substr:
                        subt = '@@@' + subname + '.srt'
                        cid = cid + subt
            videos.append({
                "vod_id":  cid,
                "vod_name": vod['name'],
                "vod_pic": img,
                "vod_tag": tag,
                "vod_remarks": remark
            })
        if 'purl' in locals().keys():
            purl = baseurl + aid + '+++' + purl
            for i in range(foldernum, lenvodList):
                if videos[i]['vod_id'] == '###':
                    videos[i]['vod_id'] = purl
        result['list'] = videos
        result['page'] = 1
        result['pagecount'] = 1
        result['limit'] = 999
        result['total'] = 999999
        return result

    def detailContent(self, array):
        id = array[0]
        if '+++' in id:
            ids = id.split('+++')
            durl = ids[0]
            vsList = ids[1].strip('#').split('#')
            vsurl = ''
            for vs in vsList:
                if '@@@' in vs:
                    dvs = vs.split('@@@')
                    vname = dvs[0]
                    vurl = durl + dvs[0]
                    surl = durl + dvs[1]
                    vsurl = vsurl + '{0}${1}@@@{2}#'.format(vname, vurl, surl)
                else:
                    vurl = durl + vs
                    vsurl = vsurl + '{0}${1}#'.format(vs, vurl)
            url = vsurl
        else:
            durl = id
        if self.ver == '' or self.baseurl == '':
            self.getVersion(durl)
        baseurl = self.baseurl
        if '+++' in id:
            vid = durl.replace(baseurl, "").strip('/')
        else:
            vid = durl.replace(re.findall(r".*/", durl)[0], "")
            url = vid + '$' + id
        vod = {
            "vod_id": vid,
            "vod_name": vid,
            "vod_pic": '',
            "vod_tag": '',
            "vod_play_from": "播放",
            "vod_play_url": url
        }
        result = {
            'list': [
                vod
            ]
        }
        return result

    def searchContent(self, key, quick):
        result = {
            'list': []
        }
        return result

    def playerContent(self, flag, id, vipFlags):
        result = {}
        url = ''
        subturl = ''
        ifsub = '@@@' in id
        if ifsub is True:
            ids = id.split('@@@')
            if self.ver == '' or self.baseurl == '':
                self.getVersion(ids[1])
            ver = self.ver
            baseurl = self.baseurl
            fileName = ids[1].replace(baseurl, "")
            vfileName = ids[0].replace(baseurl, "")
            param = {
                "path": '/' + fileName,
                "password": "",
                "page_num": 1,
                "page_size": 100
            }
            vparam = {
                "path": '/' + vfileName,
                "password": "",
                "page_num": 1,
                "page_size": 100
            }
            if ver == 2:
                rsp = self.postJson(baseurl + 'api/public/path', param)
                jo = json.loads(rsp.text)
                vodList = jo['data']['files'][0]
                subturl = vodList['url']
                vrsp = self.postJson(baseurl + 'api/public/path', vparam)
                vjo = json.loads(vrsp.text)
                vList = vjo['data']['files'][0]
                url = vList['url']
            elif ver == 3:
                rsp = self.postJson(baseurl + 'api/fs/get', param)
                jo = json.loads(rsp.text)
                vodList = jo['data']
                subturl = vodList['raw_url']
                vrsp = self.postJson(baseurl + 'api/fs/get', vparam)
                vjo = json.loads(vrsp.text)
                vList = vjo['data']
                url = vList['raw_url']
            if subturl.startswith('http') is False:
                head = re.findall(r"h.*?:", baseurl)[0]
                subturl = head + subturl
            if url.startswith('http') is False:
                head = re.findall(r"h.*?:", baseurl)[0]
                url = head + url
            urlfileName = urllib.parse.quote(fileName)
            subturl = subturl.replace(fileName, urlfileName)
            urlvfileName = urllib.parse.quote(vfileName)
            url = url.replace(vfileName, urlvfileName)
            result['subt'] = subturl
        else:
            if self.ver == '' or self.baseurl == '':
                self.getVersion(id)
            ver = self.ver
            baseurl = self.baseurl
            vfileName = id.replace(baseurl, "")
            vparam = {
                "path": '/' + vfileName,
                "password": "",
                "page_num": 1,
                "page_size": 100
            }
            if ver == 2:
                vrsp = self.postJson(baseurl + 'api/public/path', vparam)
                vjo = json.loads(vrsp.text)
                vList = vjo['data']['files'][0]
                driver = vList['driver']
                url = vList['url']
            elif ver == 3:
                vrsp = self.postJson(baseurl + 'api/fs/get', vparam)
                vjo = json.loads(vrsp.text)
                vList = vjo['data']
                url = vList['raw_url']
                driver = vList['provider']
            if url.startswith('http') is False:
                head = re.findall(r"h.*?:", baseurl)[0]
                url = head + url
            urlvfileName = urllib.parse.quote(vfileName)
            url = url.replace(vfileName, urlvfileName)
            if driver == 'Baidu.Disk':
                result["header"] = {"User-Agent": "pan.baidu.com"}
        result["parse"] = 0
        result["playUrl"] = ''
        result["url"] = url

        return result

    config = {
        "player": {},
        "filter": {}
    }
    header = {}

    def localProxy(self, param):
        return [200, "video/MP2T", action, ""]


from datasets import GeneratorBasedBuilder, BuilderConfig, Version, DatasetInfo, Features, Value, \
    Sequence, ClassLabel, DownloadManager, SplitGenerator, Split
import os
import textwrap

_CITATION = """
@inproceedings{DBLP:conf/tsd/SimonV21,
  author    = {Eszter Simon and
               No{\'{e}}mi Vad{\'{a}}sz},
  editor    = {Kamil Ekstein and
               Frantisek P{\'{a}}rtl and
               Miloslav Konop{\'{\i}}k},
  title     = {Introducing {NYTK-NerKor}, A Gold Standard {Hungarian} Named Entity
               Annotated Corpus},
  booktitle = {Text, Speech, and Dialogue - 24th International Conference, {TSD}
               2021, Olomouc, Czech Republic, September 6-9, 2021, Proceedings},
  series    = {Lecture Notes in Computer Science},
  volume    = {12848},
  pages     = {222--234},
  publisher = {Springer},
  year      = {2021},
  doi       = {10.1007/978-3-030-83527-9\_19},
}
"""
_SOURCE = "https://github.com/nytud/NYTK-NerKor/archive/36a9aefd37e1a77fb8671375def0a5ad343d5dc3.zip"
_PATHS = {
    "train": "data/train-devel-test/train/",
    "validation": "data/train-devel-test/devel/",
    "test": "data/train-devel-test/test/",
}

_SUBS = ["fiction", "legal", "news", "web", "wikipedia", "all"]

_NER = Sequence(
    ClassLabel(
        names=[
            "O",
            "B-PER",
            "I-PER",
            "B-ORG",
            "I-ORG",
            "B-LOC",
            "I-LOC",
            "B-MISC",
            "I-MISC",
        ]
    )
)

_UPOS = Sequence(
    ClassLabel(
        names=[
            "ADJ",
            "ADP",
            "ADV",
            "AUX",
            "CCONJ",
            "DET",
            "INTJ",
            "NOUN",
            "NUM",
            "PART",
            "PRON",
            "PROPN",
            "PUNCT",
            "SCONJ",
            "SYM",
            "VERB",
            "X"
        ]
    )
)

def get_repo_url(x):
    if x.startswith("https://huggingface"):
        return "/".join(x.split("/")[:6])
    return "/".join(x.split("/")[:5])

class NerKorConfig(BuilderConfig):
    """BuilderConfig for Hulu."""

    def __init__(
            self,
            text_features,
            data_url,
            data_dir,
            url,
            **kwargs,
    ):
        super(NerKorConfig, self).__init__(version=Version("1.0.0", ""), **kwargs)
        self.text_features = text_features
        self.data_url = data_url
        self.data_dir = data_dir
        self.url = url

class NerKor(GeneratorBasedBuilder):
    """NerKor datasets."""
    # FORM LEMMA UPOS XPOS FEATS CONLL:NER
    BUILDER_CONFIGS = [
        NerKorConfig(
            name=name,
            description=textwrap.dedent(""""""),
            text_features=["form", "lemma", "upos", "xpos", "feats", "ner"],
            data_url=_SOURCE,
            data_dir=f"nytud-nerkor/",
            url=get_repo_url(_SOURCE)
        ) for name in _SUBS
    ]

    def _info(self):
        features = {
            "idx": Value("int32"),
            "text": Value("string"),
            "tokens": Sequence(Value("string")),
            "lemmas": Sequence(Value("string")),
            "upos": _UPOS,
            "xpos": Sequence(Value("string")),
            "feats": Sequence(Value("string")),
            "ner": _NER,
            "file_name": Value("string"),
            "sentence_id": Value("int32"),
            "morph_tagged": Value("bool_"),
        }
        return DatasetInfo(
            description=self.config.description,
            features=Features(features),
            homepage=self.config.url,
            citation=textwrap.dedent(_CITATION),
        )

    def _split_generators(self, dl_manager: DownloadManager):
        path = dl_manager.download_and_extract(self.config.data_url)
        return [
            SplitGenerator(
                name=Split.TRAIN,
                gen_kwargs={"split_key": "train", "data_file": path},
            ),
            SplitGenerator(
                name=Split.VALIDATION,
                gen_kwargs={"split_key": "validation", "data_file": path},
            ),
            SplitGenerator(
                name=Split.TEST,
                gen_kwargs={"split_key": "test", "data_file": path},
            )
        ]

    def _generate_examples(self, data_file, split_key, **kwargs):
        if self.config.name == "all":
            base_path = os.path.join(data_file, os.listdir(data_file)[0], _PATHS[split_key])
            base_path = [os.path.join(base_path, p) for p in os.listdir(base_path)]
            train_pointers = [os.path.join(b, p) for b in base_path for p in os.listdir(b)]
        else:
            base_path = os.path.join(data_file, os.listdir(data_file)[0], _PATHS[split_key], self.config.name)
            train_pointers = [os.path.join(base_path, x) for x in os.listdir(base_path)]
        n = 0
        for annotation in train_pointers:
            for file in os.listdir(annotation):
                pointer_file_path = os.path.join(annotation, file)
                with open(pointer_file_path, mode="r") as f:
                    data_file_path = f.readlines()[0]
                for data in self._process_files(os.path.join(annotation, data_file_path), n):
                    yield data
                    n += 1

    @staticmethod
    def _process_files(data_file_path: str, n: int):
        with open(data_file_path, mode="r", encoding="utf8") as f:
            content = f.readlines()

        tokens, lemmas, upos, xpos, feats, ner = [], [], [], [], [], []

        sentence_id = 0
        for i, row in enumerate(content[1:]):
            if row == '\n':
                p = data_file_path.split("/")[-3:]
                yield n, {
                    "idx": n,
                    "text": " ".join(tokens),
                    "tokens": tokens,
                    "lemmas": lemmas,
                    "upos": upos,
                    "xpos": xpos,
                    "feats": feats,
                    "ner": ner,
                    "file_name": "/".join(p),
                    "sentence_id": sentence_id,
                    "morph_tagged": p[1] == 'morph'
                }
                tokens, lemmas, upos, xpos, feats, ner = [], [], [], [], [], []
                sentence_id += 1
                n += 1
                continue

            data = row.split("\t")
            _upos = data[2] if data[2] != "_" else -1
            _xpos = data[3] if data[3] != "_" else -1
            _feats = data[4] if data[4] != "_" else -1
            _ner = data[5] if data[5] != "_" else -1

            tokens.append(data[0])
            lemmas.append(data[1])
            upos.append(_upos)
            xpos.append(_xpos)
            feats.append(_feats)
            ner.append(_ner)


import maya.OpenMaya as om
import maya.cmds as cmds

def zbw_FK2IKSnap(*args):
    """
    select FK wrist joint(or control), fk elbow joint (or ctrl), FK shoulder joint (or ctrl), IK wrist ctl and IK pole vector in that order
    """

    sel = cmds.ls(sl=True)
    fkEndC = cmds.listConnections((sel[0]+".fkEndCtrl"))
    fkMidC = cmds.listConnections((sel[0]+".fkMidCtrl"))
    fkTopC = cmds.listConnections((sel[0]+".fkTopCtrl"))
    ikTopJ = cmds.listConnections((sel[0]+".ikTopJnt"))
    ikMidJ = cmds.listConnections((sel[0]+".ikMidJnt"))
    ikEndJ = cmds.listConnections((sel[0]+".ikEndJnt"))

    #get FK wrist joint position & rot
    ikEndPos = cmds.xform(ikEndJ, q=True, ws=True, t=True)
    iep = om.MVector(ikEndPos[0], ikEndPos[1], ikEndPos[2])
    ikEndRot = cmds.xform(ikEndJ, q=True, ro=True)
    ier = om.MVector(ikEndRot[0], ikEndRot[1], ikEndRot[2])

    #get FK shoulder position & rot
    ikTopPos = cmds.xform(ikTopJ, q=True, ws=True, t=True)
    itp = om.MVector(ikTopPos[0], ikTopPos[1], ikTopPos[2])
    ikTopRot = cmds.xform(ikTopJ, q=True, ro=True)
    itr = om.MVector(ikTopRot[0], ikTopRot[1], ikTopRot[2])

    #get midJnt pos & rot
    ikMidPos = cmds.xform(ikMidJ, q=True, ws=True, t=True)
    imp = om.MVector(ikMidPos[0], ikMidPos[1], ikMidPos[2])
    ikMidRot = cmds.xform(ikMidJ, q=True, ro=True)
    imr = om.MVector(ikMidRot[0], ikMidRot[1], ikMidRot[2])

    #rotate joints to rotations
    cmds.xform(fkEndC, ro=(ier.x, ier.y, ier.z))
    cmds.xform(fkMidC, ro=(imr.x, imr.y, imr.z))
    cmds.xform(fkTopC, ro=(itr.x, itr.y, itr.z))

    #snap FK ctrl to positions
#-------try:except this part . . . Will only work if the channels aren't locked . . .
    cmds.move(iep.x, iep.y, iep.z, fkEndC, ws=True)
    cmds.move(imp.x, imp.y, imp.z, fkMidC, ws=True)
    cmds.move(iep.x, iep.y, iep.z, fkEndC, ws=True)


import datetime
from packages.connection import *
from app.models.WhatTable import WhatTable
from app.models.WhereTable import WhereTable
from app.controllers.WhatUploader import WhatUploader
from app.controllers.WhereUploader import WhereUploader
from app.controllers.WhatController import WhatController
from app.controllers.WhereController import WhereController
from app.controllers.ResultFile import ResultFile

def showTime():
    now = datetime.datetime.now()
    result = now.strftime("%d.%m.%Y %H:%M:%S")
    return result

def showInfo(text):
    result = showTime() + ' - ' + text
    return result

def printInfo(text):
    info_text = showInfo(text)
    print(info_text)

if(__name__ == "__main__"):
    printInfo('Программа запущена!!!')
    try:
        db.connect()

        # Удаляем таблицы, если вдруг решили поменять структуру БД
        WhatTable.drop_table()
        WhereTable.drop_table()
        # Удаляем таблицы, если вдруг решили поменять структуру БД

        # Создаём таблицы
        WhatTable.create_table()
        WhereTable.create_table()
        # Создаём таблицы

        # Очищаем таблицы
        WhatTable.truncate_table()
        WhereTable.truncate_table()
        # Очищаем таблицы
    except:
        printInfo('Ошибка!!! Не удалось подключиться к базе данных!')
        exit()

    # Загружаем данные в базу
    printInfo('Данные из xlsx таблиц загружаются в базу')
    WhatUploader().run()
    WhereUploader().run()
    printInfo('Данные из xlsx таблиц успешно загружены')
    # Загружаем данные в базу

    # Создаём объект файла с результатом
    resultObj = ResultFile()
    resultObj.headerContentPut()
    # Создаём объект файла с результатом

    # Вытаскиваем все данные из БД
    what_rows = WhatController().getData()
    count_rows = len(what_rows)
    # Вытаскиваем все данные из БД

    # Перебираем содержимое БД
    printInfo('Начинаем сравнение таблиц')
    this_row = 1
    for what_row in what_rows:
        printInfo(what_row.file_name + ': Обаботана ' + str(this_row) + ' строка из ' + str(count_rows))
        this_row+=1

        where_has = WhereController().hasSnils(what_row.snils)
        if where_has['status']:
            info = {
                "what_file": what_row.file_name,
                "what_row": what_row.num_row,
                "where_file": where_has['result'].file_name,
                "where_row": where_has['result'].num_row,
                "value": what_row.snils,
                "request_date": what_row.request_date,
                "serviceid": what_row.serviceid
            }
            resultObj.contentRowPut(info)
        else:
            info = {
                "what_file": what_row.file_name,
                "what_row": what_row.num_row,
                "where_file": 'None',
                "where_row": 'None',
                "value": what_row.snils,
                "request_date": what_row.request_date,
                "serviceid": what_row.serviceid
            }
            resultObj.contentRowPut(info)
    printInfo('Сравнение таблиц успешно завершено')

    # Сохраняем файл результата
    printInfo('Сохраняем результа в файл')
    resultObj.saveFile()
    # Сохраняем файл результата

    printInfo('Программа выполнена!!!')
    print('')
    print('***')
    print('Название программы: checkSnils v3.0')
    print('Разработка: Министерства труда и социального развития РД')
    print('Разработчик: Ахмедов Мурад Алилович')
    print('***')
    print('')
    input('Нажмите любую клавишу, для завершения работы программы')


import os

import jax.numpy as jnp
import numpy as np
from experiment_defs import *

def split_merge_vary_num_dataset(dir_path, group_id, configs):
    experiment_name = 'hpl_bo_split_group_id_{}_merge'.format(group_id)
    results = {}

    results['experiment_name'] = experiment_name

    # configs
    results['configs'] = configs
    results['configs']['hpl_lengthscale_mlp_features'] = HPL_LENGTHSCALE_MLP_FEATURES
    results['configs']['vary_num_dataset_method_name_list'] = VARY_NUM_DATASET_METHOD_NAME_LIST

    setup_b_id_list = configs['setup_b_id_list']
    method_name_list = configs['method_name_list']
    ac_func_type_list = configs['ac_func_type_list']
    n_bo_runs = configs['n_bo_runs']
    eval_nll_n_batches = configs['eval_nll_n_batches']

    results['setup_b'] = {}

    for num_dataset in VARY_NUM_DATASET_NUM_LIST:
        for seed in range(VARY_NUM_DATASET_NUM_SEEDS):
            extra_suffix = '_n{}seed{}'.format(num_dataset, seed)
            # setup b
            results['setup_b'][extra_suffix] = {}
            results_b = results['setup_b'][extra_suffix]

            # fit gp parameters
            fit_single_gp_list = setup_b_id_list
            results_b['fit_base_gp_params'] = {}
            for dataset_id in fit_single_gp_list:
                results_b['fit_base_gp_params'][dataset_id] = np.load(
                    os.path.join(dir_path,
                                 'split_fit_gp_params_setup_b_id_{}.npy'.format(dataset_id)), allow_pickle=True).item()

            # read fit direct hgp params
            if 'fit_direct_hgp' in VARY_NUM_DATASET_METHOD_NAME_LIST:
                results_b['fit_direct_hgp_params'] = np.load(
                    os.path.join(dir_path, 'split_fit_direct_hgp_two_step_setup_b{}.npy'.format(extra_suffix)),
                    allow_pickle=True).item()['gp_distribution_params']

            # read fit hpl hgp params
            if 'hpl_hgp_end_to_end' in VARY_NUM_DATASET_METHOD_NAME_LIST:
                results_b['hpl_hgp_end_to_end_params'] = np.load(
                    os.path.join(dir_path, 'split_fit_hpl_hgp_end_to_end_setup_b{}.npy'.format(extra_suffix)),
                    allow_pickle=True).item()
            if 'hpl_hgp_two_step' in VARY_NUM_DATASET_METHOD_NAME_LIST:
                results_b['hpl_hgp_two_step_params'] = np.load(
                    os.path.join(dir_path, 'split_fit_hpl_hgp_two_step_setup_b{}.npy'.format(extra_suffix)),
                    allow_pickle=True).item()

            # run BO and compute NLL
            results_b['bo_results'] = {}
            results_b['bo_results_total'] = {}
            results_b['nll_results'] = {}
            results_b['nll_results_total'] = {}
            results_b['kl_eval'] = {}

            for method_name in VARY_NUM_DATASET_METHOD_NAME_LIST:
                # BO results
                results_b['bo_results'][method_name] = {}
                results_b['bo_results_total'][method_name] = {}

                for ac_func_type in ac_func_type_list:
                    results_b['bo_results'][method_name][ac_func_type] = {}

                    regrets_all_list = []

                    for test_id in TEST_ID_LIST:
                        results_b['bo_results'][method_name][ac_func_type][test_id] = np.load(
                            os.path.join(dir_path,
                                         'split_test_bo_setup_b_id_{}_{}_{}{}.npy'.format(test_id, method_name,
                                                                                ac_func_type,
                                                                                extra_suffix)),
                            allow_pickle=True,
                        ).item()

                    for i in range(n_bo_runs):
                        regrets_all_list_i = []

                        for test_id in TEST_ID_LIST:
                            regrets_ij = results_b['bo_results'][method_name][ac_func_type][test_id]['regrets_list'][i]
                            regrets_all_list_i.append(regrets_ij)

                        regrets_all_list.append(jnp.mean(jnp.array(regrets_all_list_i), axis=0))

                    regrets_all_list = jnp.array(regrets_all_list)
                    regrets_mean_total = jnp.mean(regrets_all_list, axis=0)
                    regrets_std_total = jnp.std(regrets_all_list, axis=0)

                    results_b['bo_results_total'][method_name][ac_func_type] = {
                        'regrets_all_list': regrets_all_list,
                        'regrets_mean': regrets_mean_total,
                        'regrets_std': regrets_std_total,
                    }

                # KL divergence evaluation
                if method_name in EVAL_KL_METHOD_NAME_LIST:
                    results_b['kl_eval'][method_name] = np.load(os.path.join(
                        dir_path, 'split_eval_kl_with_gt_{}{}.npy'.format(method_name, extra_suffix)),
                        allow_pickle=True).item()

    # save all results
    merge_path = os.path.join(dir_path, 'merge')
    if not os.path.exists(merge_path):
        os.mkdir(merge_path)
    np.save(os.path.join(merge_path, 'results.npy'), results)

    # write part of results to text file
    with open(os.path.join(merge_path, 'results.txt'), 'w') as f:
        f.write(str(results))

    print('done.')