"""
Provide registration functions for numerical grid parameters.

This module handles grid parameters and grid functions across various NRPy-supported infrastructures.

Author: Zachariah B. Etienne
        zachetie **at** gmail **dot* com
"""
from typing import List, Any, Union, Dict, Tuple, Optional, Sequence
import sympy as sp
import nrpy.indexedexp as ixp
import nrpy.params as par

Cart_origin = par.register_CodeParameters(
    "REAL", __name__, ["Cart_originx", "Cart_originy", "Cart_originz"], 0.0
)
_ = par.register_CodeParameter("int", __name__, "NUMGRIDS", 1, commondata=True)

# Core GridFunction class.
class GridFunction:
    """The core class for grid functions."""

    def __init__(self, name: str, group: str = "EVOL") -> None:
        self.c_type_alias: str = "double"
        self.name: str = name
        self.group: str = group

    def read_gf_from_memory_Ccode_onept(
        self, i0_offset: int = 0, i1_offset: int = 0, i2_offset: int = 0, **kwargs: Any
    ) -> str:
        """Catch-all function for gridfunctions that haven't been set up correctly."""
        return f"Please define read_gf_from_memory_Ccode_onept() inside grid.py. Inputs: [class instance] {i0_offset} {i1_offset} {i2_offset} {kwargs}"

class BHaHGridFunction(GridFunction):
    """The subclass for BlackHoles@Home grid functions."""

    def __init__(
        self,
        name: str,
        group: str = "EVOL",
        rank: int = 0,
        f_infinity: float = 0.0,
        wavespeed: float = 1.0,
        is_basename: bool = True,
        gf_array_name: str = "use_in_gfs_for_EVOL_auxevol_gfs_for_AUXEVOL_etc",
    ) -> None:
        super().__init__(name, group)
        self.c_type_alias = "REAL"  # always use REAL
        self.rank = rank
        self.f_infinity = f_infinity
        self.wavespeed = wavespeed
        self.is_basename = is_basename
        self.verify_gridfunction_group_is_valid()
        if gf_array_name == "use_in_gfs_for_EVOL_auxevol_gfs_for_AUXEVOL_etc":
            if group == "EVOL":
                self.gf_array_name = "in_gfs"
            elif group == "AUXEVOL":
                self.gf_array_name = "auxevol_gfs"
            elif group == "AUX":
                self.gf_array_name = "aux_gfs"
        else:
            self.gf_array_name = gf_array_name

        if is_basename:
            self.verify_gridfunction_basename_is_valid(self.name)

    def verify_gridfunction_basename_is_valid(self, name: str) -> None:
        """
        Validate the gridfunction's base name.

        Raises a ValueError if base name is zero-length or ends with an integer.
        Raises TypeError if the base name is not a string.

        Gridfunction base names must not end with integers. For instance, a rank-1
        gridfunction 'vecU1' has a valid base name 'vecU'. A scalar gridfunction
        named 'u2' would have an invalid base name 'u2'. This strict requirement
        facilitates quick identification of a gridfunction by its name.

        :param name: The name of the gridfunction.

        Raises
        ------
        ValueError
            If the name is of zero length or ends with an integer.
        TypeError
            If the name is not of type str.

        Example:
            >>> try: GF = BHaHGridFunction(name="h1")
            ... except: "Attempt to create a gridfunction ending in an integer failed. Good."
            'Attempt to create a gridfunction ending in an integer failed. Good.'
        """
        # First check for zero-length basenames:
        if len(name) == 0:
            raise ValueError("Tried to register gridfunction without a name!")

        if not isinstance(name, str):
            raise TypeError("Gridfunction names must be strings")

        if len(name) > 0 and name[-1].isdigit():
            raise ValueError(
                f"Tried to register gridfunction with base name: {name}\n"
                f"To ensure that finite difference code generations do not get confused,\n"
                f"gridfunctions with base names ending in an integer are forbidden; pick a new name."
            )

    def verify_gridfunction_group_is_valid(self) -> None:
        """
        Validate the gridfunction group.

        The valid groups are 'EVOL', 'AUX' and 'AUXEVOL'.

        'EVOL': for evolved quantities (i.e., quantities stepped forward in time),
        'AUXEVOL': for auxiliary quantities needed at all points by evolved quantities,
        'AUX': for all other quantities needed at all gridpoints.

        :raises ValueError: If the group is not valid.

        Doctests:
        >>> try: GF = BHaHGridFunction(name="blah", group="INVALID")
        ... except: print("Errored out. This is good.")
        Errored out. This is good.
        """
        valid_groups = (
            "EVOL",
            "AUX",
            "AUXEVOL",
        )
        if self.group not in valid_groups:
            msg = (
                f"Unsupported gridfunction group {self.group}. Supported groups include:\n"
                '    "EVOL": for evolved quantities (i.e., quantities stepped forward in time),\n'
                '    "AUXEVOL": for auxiliary quantities needed at all points by evolved quantities,\n'
                '    "AUX": for all other quantities needed at all gridpoints.\n'
            )
            raise ValueError(msg)

    def read_gf_from_memory_Ccode_onept(
        self, i0_offset: int = 0, i1_offset: int = 0, i2_offset: int = 0, **kwargs: Any
    ) -> str:
        """
        Retrieve a grid function value from memory for a given offset.

        :param i0_offset: Offset for the first index.
        :param i1_offset: Offset for the second index.
        :param i2_offset: Offset for the third index.
        :param kwargs: Optional keyword arguments, expected to have 'gf_array_name'.

        :return: Formatted string.

        :raises ValueError: If 'gf_array_name' is not provided in kwargs.

        Doctests:
        >>> glb_gridfcs_dict.clear()
        >>> par.set_parval_from_str("Infrastructure", "BHaH")
        >>> abc = register_gridfunctions("abc")
        >>> glb_gridfcs_dict["abc"].read_gf_from_memory_Ccode_onept(1, 2, 3)
        'in_gfs[IDX4(ABCGF, i0+1, i1+2, i2+3)]'
        >>> defg = register_gridfunctions("defg")
        >>> glb_gridfcs_dict["defg"].read_gf_from_memory_Ccode_onept(0, -1, 0, enable_simd=True, gf_array_name="My_Array")
        'ReadSIMD(&My_Array[IDX4(DEFGGF, i0, i1-1, i2)])'
        """
        i0 = f"i0+{i0_offset}".replace("+-", "-") if i0_offset != 0 else "i0"
        i1 = f"i1+{i1_offset}".replace("+-", "-") if i1_offset != 0 else "i1"
        i2 = f"i2+{i2_offset}".replace("+-", "-") if i2_offset != 0 else "i2"

        gf_array_name = kwargs.get("gf_array_name", self.gf_array_name)

        ret_string = f"{gf_array_name}[IDX4({self.name.upper()}GF, {i0}, {i1}, {i2})]"
        if kwargs.get("enable_simd"):
            return f"ReadSIMD(&{ret_string})"
        return ret_string

    @staticmethod
    def access_gf(
        gf_name: str,
        i0_offset: int = 0,
        i1_offset: int = 0,
        i2_offset: int = 0,
        gf_array_name: str = "in_gfs",
    ) -> str:
        """
        Retrieve a grid function value from memory for a given offset.

        :param gf_name: The grid function name.
        :param i0_offset: Offset for the first index.
        :param i1_offset: Offset for the second index.
        :param i2_offset: Offset for the third index.
        :param gf_array_name: Optional grid function array name.

        :return: Formatted string.

        :raises ValueError: If 'gf_array_name' is not provided.

        Doctests:
        >>> BHaHGridFunction.access_gf("aa", 1,2,3)
        'in_gfs[IDX4(AAGF, i0+1, i1+2, i2+3)]'
        >>> BHaHGridFunction.access_gf("defg", 0, -1, 0, "My_Array")
        'My_Array[IDX4(DEFGGF, i0, i1-1, i2)]'
        """
        i0 = f"i0+{i0_offset}".replace("+-", "-") if i0_offset != 0 else "i0"
        i1 = f"i1+{i1_offset}".replace("+-", "-") if i1_offset != 0 else "i1"
        i2 = f"i2+{i2_offset}".replace("+-", "-") if i2_offset != 0 else "i2"

        return f"{gf_array_name}[IDX4({gf_name.upper()}GF, {i0}, {i1}, {i2})]"

    @staticmethod
    def gridfunction_lists() -> Tuple[List[str], List[str], List[str]]:
        """
        Generate sorted lists of gridfunction names for each group type: 'EVOL', 'AUX', and 'AUXEVOL'.

        This function creates a dictionary with group types as keys and corresponding gridfunction names
        as values. It iterates through a global dictionary of gridfunctions, 'glb_gridfcs_dict', and
        appends the gridfunction names to the corresponding group lists in 'groups'. The function then
        sorts each list and returns the sorted lists as a tuple.

        Note: This function assumes the existence of a global gridfunction dictionary named
        'glb_gridfcs_dict' whose keys are gridfunction names and values are gridfunction objects. Each
        gridfunction object should have 'group' and 'name' attributes.

        :returns: A tuple containing lists of gridfunction names
                  for each group: 'EVOL', 'AUX', and 'AUXEVOL', respectively.

        Example:
            gridfunction_lists()
            (['evol_gf1', 'evol_gf2'], ['aux_gf1', 'aux_gf2'], ['auxevol_gf1', 'auxevol_gf2'])
        """
        # Initialize dictionary for holding lists of gridfunction names for each group.
        groups: Dict[str, List[str]] = {
            "EVOL": [],
            "AUX": [],
            "AUXEVOL": [],
        }

        # Iterate through the global dictionary of gridfunctions.
        for _key, gf in glb_gridfcs_dict.items():
            if gf.group in groups:
                groups[gf.group].append(gf.name)

        # Sort the lists. Iterating through a copy of the keys to avoid modifying the dictionary while iterating.
        for group in list(groups.keys()):
            groups[group] = sorted(groups[group])

        # Pack the sorted lists into a tuple and return.
        return groups["EVOL"], groups["AUX"], groups["AUXEVOL"]

    @staticmethod
    def gridfunction_defines() -> str:
        """
        Generate a string representation of all the GridFunction definitions.

        This function also includes definitions for the values of each gridfunction
        at infinity and the characteristic wave speed. The output is formatted to be
        used in C-style languages.

        This function assumes the existence of a global gridfunction dictionary named
        'glb_gridfcs_dict'. The keys of this dictionary should be gridfunction names,
        and the values should be instances of classes derived from GridFunction.

        :return: The resulting string representation of the GridFunction definitions.
        """

        def define_gfs(name: str, gfs: List[str]) -> str:
            """
            Generate string representation for a list of GridFunctions.

            :param name: The name of the group.
            :param gfs: The list of GridFunctions.
            :return: The resulting string representation.
            """
            num_gfs = len(gfs)
            defines = "\n".join(
                f"#define {gf.upper()}GF\t{i}" for i, gf in enumerate(gfs)
            )
            return f"\n// {name.upper()} VARIABLES:\n#define NUM_{name.upper()}_GFS {num_gfs}\n{defines}\n"

        (
            evolved_variables_list,
            auxiliary_variables_list,
            auxevol_variables_list,
        ) = BHaHGridFunction.gridfunction_lists()

        outstr = f"{define_gfs('EVOL', evolved_variables_list)}"

        if evolved_variables_list:
            # mypy: doesn't consider that this func is only called for BHaHGridFunctions.
            f_infinity_list = [str(glb_gridfcs_dict[var].f_infinity) for var in evolved_variables_list]  # type: ignore
            f_wavespeed_list = [str(glb_gridfcs_dict[var].wavespeed) for var in evolved_variables_list]  # type: ignore

            f_infinity_str = ", ".join(f_infinity_list)
            f_wavespeed_str = ", ".join(f_wavespeed_list)

            outstr += "\n\n// SET gridfunctions_f_infinity[i] = evolved gridfunction i's value in the limit r->infinity:\n"
            outstr += f"static const REAL gridfunctions_f_infinity[NUM_EVOL_GFS] = {{ {f_infinity_str} }};\n"

            outstr += "\n\n// SET gridfunctions_wavespeed[i] = evolved gridfunction i's characteristic wave speed:\n"
            outstr += f"static const REAL gridfunctions_wavespeed[NUM_EVOL_GFS] = {{ {f_wavespeed_str} }};\n"

        outstr += f"{define_gfs('AUX', auxiliary_variables_list)}"
        outstr += f"{define_gfs('AUXEVOL', auxevol_variables_list)}"

        return outstr

class BaseETGridFunction(GridFunction):
    """Subclass for basic (non-CarpetX) Einstein Toolkit grid functions."""

    def __init__(
        self,
        name: str,
        group: str = "EVOL",
        rank: int = 0,
        f_infinity: float = 0.0,
        wavespeed: float = 1.0,
        is_basename: bool = True,
    ) -> None:
        super().__init__(name, group)
        self.rank = rank
        self.f_infinity = f_infinity
        self.wavespeed = wavespeed
        self.is_basename = is_basename
        self.c_type_alias = "CCTK_REAL"

    def verify_gridfunction_group_is_valid(self) -> None:
        """
        Verify that the 'group' attribute of the GridFunction instance is one of the valid groups.

        The valid groups are:
        'EVOL': for evolved quantities (i.e., quantities stepped forward in time),
        'AUXEVOL': for auxiliary quantities needed at all points by evolved quantities,
        'AUX': for all other quantities needed at all gridpoints.

        :raises ValueError: If the 'group' attribute is not one of the valid groups.
        """
        valid_groups = (
            "EVOL",
            "AUX",
            "AUXEVOL",
        )
        if self.group not in valid_groups:
            msg = (
                f"Unsupported gridfunction group {self.group}. Supported groups include:\n"
                '    "EVOL": for evolved quantities (i.e., quantities stepped forward in time),\n'
                '    "AUXEVOL": for auxiliary quantities needed at all points by evolved quantities,\n'
                '    "AUX": for all other quantities needed at all gridpoints.\n'
            )
            raise ValueError(msg)

    def read_gf_from_memory_Ccode_onept(
        self, i0_offset: int = 0, i1_offset: int = 0, i2_offset: int = 0, **kwargs: Any
    ) -> str:
        """
        Generate a formatted string using the grid function name and offsets.

        If 'gf_array_name' is not provided in kwargs, a ValueError is raised.

        :param i0_offset: Offset for the first index
        :param i1_offset: Offset for the second index
        :param i2_offset: Offset for the third index
        :param kwargs: Optional keyword arguments.
        :return: Formatted string
        :raises ValueError: If 'gf_array_name' is not provided in kwargs or if kwargs are passed when they shouldn't be.

        Doctests:
        >>> glb_gridfcs_dict.clear()
        >>> par.set_parval_from_str("Infrastructure", "BaseET")
        >>> abc = register_gridfunctions("abc", group="EVOL")
        >>> glb_gridfcs_dict["abc"].read_gf_from_memory_Ccode_onept(1, 2, 3)
        'abc[CCTK_GFINDEX3D(cctkGH, i0+1, i1+2, i2+3)]'
        >>> defg = register_gridfunctions("defg", group="EVOL")
        >>> glb_gridfcs_dict["defg"].read_gf_from_memory_Ccode_onept(0, -1, 0)
        'defg[CCTK_GFINDEX3D(cctkGH, i0, i1-1, i2)]'
        """
        if kwargs:
            raise ValueError(
                "BaseETGridFunction.read_gf_from_memory_Ccode_onept() does not accept kwargs!"
            )

        i0 = f"i0+{i0_offset}".replace("+-", "-") if i0_offset != 0 else "i0"
        i1 = f"i1+{i1_offset}".replace("+-", "-") if i1_offset != 0 else "i1"
        i2 = f"i2+{i2_offset}".replace("+-", "-") if i2_offset != 0 else "i2"

        return f"{self.name}[CCTK_GFINDEX3D(cctkGH, {i0}, {i1}, {i2})]"

class CarpetXGridFunction(GridFunction):
    """Class for CarpetX gridfunction operations."""

    def __init__(
        self,
        name: str,
        group: str = "",
        rank: int = 0,
        f_infinity: float = 0.0,
        wavespeed: float = 1.0,
        is_basename: bool = True,
        centering: str = "C",
    ) -> None:
        super().__init__(name, group)
        self.rank = rank
        self.f_infinity = f_infinity
        self.wavespeed = wavespeed
        self.is_basename = is_basename
        self.centering = centering

        group_suffixes = {
            "EXTERNAL": "_ext",
            "CORE": "_core",
            "TILE_TMP": "_tile_tmp",
            "SCALAR_TMP": "",
        }
        if self.group in group_suffixes:
            self.name += group_suffixes[self.group]

        # Dictionaries for storing grid function groups.
        # self.index_group = {}
        # self.rev_index_group = {}

        # Validation step.
        self.verify_gridfunction_centering_is_valid()
        self.verify_gridfunction_group_is_valid()

    def verify_gridfunction_group_is_valid(self) -> None:
        """
        Verify that the 'group' attribute of the GridFunction instance is one of the valid groups.

        The valid groups are:
        'EVOL': for evolved quantities (i.e., quantities stepped forward in time),
        'AUXEVOL': for auxiliary quantities needed at all points by evolved quantities,
        'AUX': for all other quantities needed at all gridpoints,
        'EXTERNAL': for all quantities defined in other modules,
        'CORE': for all quantities defined inside the CarpetX driver,
        'TILE_TMP': for all temporary quantities defined for CarpetX tiles,
        'SCALAR_TMP': for all temporary quantities defined for doubles.

        :raises ValueError: If the 'group' attribute is not one of the valid groups.
        """
        valid_groups = (
            "EVOL",
            "AUX",
            "AUXEVOL",
            "EXTERNAL",
            "CORE",
            "TILE_TMP",
            "SCALAR_TMP",
        )
        if self.group not in valid_groups:
            msg = (
                f"Unsupported gridfunction group {self.group}. Supported groups include:\n"
                '    "EVOL": for evolved quantities (i.e., quantities stepped forward in time),\n'
                '    "AUXEVOL": for auxiliary quantities needed at all points by evolved quantities,\n'
                '    "AUX": for all other quantities needed at all gridpoints.\n'
                '    "EXTERNAL": for all quantities defined in other modules.\n'
                '    "CORE": for all quantities defined inside the CarpetX driver.\n'
                '    "TILE_TMP": for all temporary quantities defined for CarpetX tiles.\n'
                '    "SCALAR_TMP": for all temporary quantities defined for doubles.'
            )
            raise ValueError(msg)

    def verify_gridfunction_centering_is_valid(self) -> None:
        """
        Verify that the 'centering' attribute of the GridFunction instance is valid.

        A valid 'centering' is a 3-character string where each character is either 'C' or 'V'.
        'C' stands for cell-centered quantities, 'V' stands for vertex-centered quantities.

        :raises ValueError: If the 'centering' attribute is not valid.
        """
        if len(self.centering) != 3 or any(c not in ("C", "V") for c in self.centering):
            msg = (
                f"Unsupported gridfunction centering {self.centering}. Supported centerings are 3-character strings, each character being either:\n"
                '    "C": for cell-centered quantities,\n'
                '    "V": for vertex-centered quantities.'
            )
            raise ValueError(msg)

# Contains a list of gridfunction objects.
glb_gridfcs_dict: Dict[
    str, Union[GridFunction, BHaHGridFunction, BaseETGridFunction, CarpetXGridFunction]
] = {}

#####################################
# Outside the classes above, only define functions in grid.py that are used by multiple infrastructures

def register_gridfunctions(
    names: Union[str, List[str]], **kwargs: Any
) -> List[sp.Symbol]:
    """
    Register grid functions with a specified name or list of names.

    :param names: A name or list of names of grid functions to be registered.
    :param kwargs: Additional keyword arguments.
    :return: A list of sympy symbols representing the registered grid functions.

    Doctests:
    >>> glb_gridfcs_dict.clear()
    >>> gridfunc = register_gridfunctions('gridfunc')[0]
    >>> print(gridfunc)
    gridfunc
    >>> # glb_gridfcs_dict["gridfunc"] is of type BHaHGridFunction. __dict__.items() access all variables in the class.
    >>> for key, value in glb_gridfcs_dict["gridfunc"].__dict__.items():
    ...     print(key, value)
    c_type_alias CCTK_REAL
    name gridfunc
    group EVOL
    rank 0
    f_infinity 0.0
    wavespeed 1.0
    is_basename True
    >>> gridfunc1, gridfunc2 = register_gridfunctions(['gridfunc1', 'gridfunc2'], f_infinity=[1.0, 4.0])
    >>> for key, value in glb_gridfcs_dict["gridfunc1"].__dict__.items():
    ...     print(key, value)
    c_type_alias CCTK_REAL
    name gridfunc1
    group EVOL
    rank 0
    f_infinity [1.0, 4.0]
    wavespeed 1.0
    is_basename True
    """
    # Step 1: Convert names to a list if it's not already a list
    if not isinstance(names, list):
        names = [names]

    # Step 2: Process each gridfunction
    sympy_symbol_list: List[sp.Symbol] = []
    Infrastructure = par.parval_from_str("Infrastructure")
    for i, name in enumerate(names):
        if name in glb_gridfcs_dict:
            print(f"Warning: Gridfunction {name} is already registered.")
        else:
            gf: Union[BHaHGridFunction, BaseETGridFunction, CarpetXGridFunction]
            if Infrastructure == "BHaH":
                kwargs_modify = kwargs.copy()
                for param in ["f_infinity", "wavespeed"]:
                    if kwargs.get(param) and isinstance(kwargs.get(param), list):
                        # mypy: Once again bonks out after I've CONFIRMED kwargs.get(param) is not None and is a list!
                        kwargs_modify[param] = kwargs.get(param)[i]  # type: ignore
                gf = BHaHGridFunction(name, **kwargs_modify)
            elif Infrastructure == "BaseET":
                gf = BaseETGridFunction(name, **kwargs)
            elif Infrastructure == "CarpetX":
                gf = CarpetXGridFunction(name, **kwargs)
            else:
                raise ValueError(f"Infrastructure = {Infrastructure} unknown")

            glb_gridfcs_dict[name] = gf
        sympy_symbol_list += [sp.symbols(name, real=True)]

    # Step 3: Return the list of sympy symbols, which can have any number of elements >= 1
    return sympy_symbol_list

def register_gridfunctions_for_single_rank1(
    basename: str, **kwargs: Any
) -> Sequence[sp.Expr]:
    """
    Register gridfunctions for a single rank 1 variable.

    :param basename: Base name of the gridfunction to register.
    :param kwargs: Additional keyword arguments for registration.
    :return: List of SymPy symbols.

    Doctests:
    >>> glb_gridfcs_dict.clear()
    >>> register_gridfunctions_for_single_rank1("betU")
    [betU0, betU1, betU2]

    Now make sure this was stored to glb_gridfcs_dict correctly:
    >>> outstr = ""
    >>> for key, gfobj in glb_gridfcs_dict.items():
    ...    outstr += gfobj.name + " "
    >>> print(outstr[:-1])
    betU0 betU1 betU2
    """
    dimension = kwargs.get("dimension", 3)
    if "dimension" in kwargs:
        del kwargs["dimension"]

    # Step 1: Declare a list of SymPy variables,
    #         where IDX_OBJ_TMP[i] = gf_basename+str(i)
    IDX_OBJ_TMP = ixp.declarerank1(basename, dimension=dimension, **kwargs)

    # Step 2: Register each gridfunction
    gf_list = []
    for i in range(3):
        gf_list += [str(IDX_OBJ_TMP[i])]
    # Manually adjust kwargs so that basename is not checked (we're not passing the basename!),
    # and rank is set to 1.
    kwargs["is_basename"] = False
    kwargs["rank"] = 1
    register_gridfunctions(gf_list, **kwargs)

    # Step 3: Return array of SymPy variables
    return IDX_OBJ_TMP

def register_gridfunctions_for_single_rank2(
    basename: str, symmetry: Optional[str] = None, **kwargs: Any
) -> Sequence[Sequence[sp.Expr]]:
    """
    Register gridfunctions for a single rank 2 variable.

    :param basename: Base name of the gridfunction to register.
    :param symmetry: Optional symmetry for rank 2 gridfunctions.
    :param kwargs: Additional keyword arguments for registration.
    :return: 2D list of SymPy symbols.

    Doctests:
    >>> glb_gridfcs_dict.clear()
    >>> register_gridfunctions_for_single_rank2("gDD", symmetry="sym01")
    [[gDD00, gDD01, gDD02], [gDD01, gDD11, gDD12], [gDD02, gDD12, gDD22]]

    Now make sure this was stored to glb_gridfcs_dict correctly:
    >>> outstr = ""
    >>> for key, gfobj in glb_gridfcs_dict.items():
    ...    outstr += gfobj.name + " "
    >>> print(outstr[:-1])
    gDD00 gDD01 gDD02 gDD11 gDD12 gDD22
    """
    dimension = kwargs.get("dimension", 3)
    if "dimension" in kwargs:
        del kwargs["dimension"]

    # Step 1: Declare a list of lists of SymPy variables,
    #         where IDX_OBJ_TMP[i][j] = gf_basename+str(i)+str(j)
    IDX_OBJ_TMP = ixp.declarerank2(
        basename, symmetry=symmetry, dimension=dimension, **kwargs
    )

    # Step 2: register each gridfunction, being careful not
    #         not to store duplicates due to rank-2 symmetries.
    gf_list: List[str] = []
    for i in range(dimension):
        for j in range(dimension):
            save = True
            for gf in gf_list:
                if gf == str(IDX_OBJ_TMP[i][j]):
                    save = False
            if save:
                gf_list.append(str(IDX_OBJ_TMP[i][j]))

    # Manually adjust kwargs so that basename is not checked (we're not passing the basename!),
    # and rank is set to 2.
    kwargs["is_basename"] = False
    kwargs["rank"] = 2
    register_gridfunctions(gf_list, **kwargs)

    # Step 3: Return array of SymPy variables
    return IDX_OBJ_TMP

# Outside the classes above, only define functions in grid.py that are used by multiple infrastructures
#####################################

if __name__ == "__main__":
    import doctest
    import sys

    results = doctest.testmod()

    if results.failed > 0:
        print(f"Doctest failed: {results.failed} of {results.attempted} test(s)")
        sys.exit(1)
    else:
        print(f"Doctest passed: All {results.attempted} test(s) passed")


import ROOT
from basf2 import B2INFO
import flavorTagger as ft
from defaultEvaluationParameters import categories, Quiet, r_subsample, r_size
from array import array
import pickle
import math
import glob
import sys

ROOT.gROOT.SetBatch(True)

if len(sys.argv) != 3:
    sys.exit("Must provide 2 arguments: [input_sim_file] or ['input_sim_file*'] with wildcards and [treeName]"
             )
workingFile = sys.argv[1]
workingFiles = glob.glob(str(workingFile))
treeName = str(sys.argv[2])

if len(workingFiles) < 1:
    sys.exit("No file name or file names " + str(workingFile) + " found.")

workingDirectory = '.'

#
# *******************************************
# DETERMINATION OF TOTAL EFFECTIVE EFFICIENCY
# *******************************************
#

# working directory
# needs the B0bar_final.root-file
# treeName = 'B0tree'

methods = []

tree = ROOT.TChain(treeName)

mcstatus = array('d', [-511.5, 0.0, 511.5])
ROOT.TH1.SetDefaultSumw2()

for iFile in workingFiles:
    tree.AddFile(iFile)

totalBranches = []
for branch in tree.GetListOfBranches():
    totalBranches.append(branch.GetName())

if 'FBDT_qrCombined' in totalBranches:
    methods.append("FBDT")

if 'FANN_qrCombined' in totalBranches:
    methods.append("FANN")

if 'DNN_qrCombined' in totalBranches:
    methods.append("DNN")

usedCategories = []
for cat in categories:
    catBranch = 'qp' + cat
    if catBranch in totalBranches:
        usedCategories.append(cat)

YmaxForQrPlot = 0

total_notTagged = 0

for method in methods:
    # Get error with GetBinError(), set error with SetBinError()
    # histogram contains the average r in each of 7 bins -> calculation see below
    histo_avr_r = ROOT.TH1F('Average_r', 'Average r in each of the bins (B0 and B0bar)', int(r_size - 2),
                            r_subsample)
    histo_avr_rB0 = ROOT.TH1F('Average_rB0', 'Average r in each of the bins (B0)', int(r_size - 2),
                              r_subsample)
    histo_avr_rB0bar = ROOT.TH1F('Average_rB0bar', 'Average r in each of the bins (B0bar)', int(r_size - 2),
                                 r_subsample)

    # histogram containing the wrong tags calculated by counting Wrong tags
    histo_mc_NwB0 = ROOT.TH1F('mc_NwB0', 'Average r in each of the bins (B0)', int(r_size - 2),
                              r_subsample)
    histo_mc_NwB0bar = ROOT.TH1F('mc_NwB0bar', 'Average r in each of the bins (B0bar)', int(r_size - 2),
                                 r_subsample)

    # histogram contains the mean squared of r in each of 7 bins -> calculation see below
    histo_ms_r = ROOT.TH1F('MS_r', 'Mean squared average of r in each of the bins (B0 and B0bar)', int(r_size - 2),
                           r_subsample)
    histo_ms_rB0 = ROOT.TH1F('MS_rB0', 'Mean squared average of r in each of the bins (B0)', int(r_size - 2),
                             r_subsample)
    histo_ms_rB0bar = ROOT.TH1F('MS_rB0bar', 'Mean squared average of r in each of the bins (B0bar)', int(r_size - 2),
                                r_subsample)

    # histogram with number of entries in for each bin
    histo_entries_per_bin = ROOT.TH1F(
        'entries_per_bin',
        'Events binned in r_subsample according to their r-value for B0 and B0bar prob',
        int(r_size - 2),
        r_subsample)
    histo_entries_per_binB0 = ROOT.TH1F('entries_per_binB0', 'Events binned in r_subsample according '
                                        'to their r-value for B0 prob', int(r_size - 2), r_subsample)
    histo_entries_per_binB0bar = ROOT.TH1F('entries_per_binB0bar',
                                           'Events binned in r_subsample according to their r-value '
                                           'for B0bar prob', int(r_size - 2), r_subsample)
    # histogram network output (not qr and not r) for true B0 (signal) - not necessary
    histo_Cnet_output_B0 = ROOT.TH1F('Comb_Net_Output_B0', 'Combiner network output [not equal to r] '
                                     'for true B0 (binning 100)', 100, 0.0, 1.0)
    # histogram network output (not qr and not r) for true B0bar (background) - not necessary
    histo_Cnet_output_B0bar = ROOT.TH1F('Comb_Net_Output_B0bar', 'Combiner network output [not equal to r] '
                                        'for true B0bar (binning 100)', 100, 0.0, 1.0)
    # histogram containing the belle paper plot (qr-tagger output for true B0)
    histo_belleplotB0 = ROOT.TH1F('BellePlot_B0',
                                  'BellePlot for true B0 (binning 50)', 50,
                                  -1.0, 1.0)
    # histogram containing the belle paper plot (qr-tagger output for true B0bar)
    histo_belleplotB0bar = ROOT.TH1F('BellePlot_B0Bar',
                                     'BellePlot for true B0Bar (binning 50)',
                                     50, -1.0, 1.0)

    histo_notTaggedEvents = ROOT.TH1F('notTaggedEvents',
                                      'Histogram for not tagged events',
                                      1, -3.0, -1.0)

    # calibration plot for B0. If we get a linaer line our MC is fine, than the assumption r ~ 1- 2w is reasonable
    # expectation is, that for B0 calibration plot:  qr=0  half B0 and half B0bar, qr = 1 only B0 and qr = -1
    # no B0. Inverse for B0bar calibration plot
    histo_calib_B0 = ROOT.TH1F('Calibration_B0', 'CalibrationPlot for true B0', 100, -1.0, 1.0)
    # calibration plot for B0bar calibration plot
    histo_calib_B0bar = ROOT.TH1F('Calibration_B0Bar',
                                  'CalibrationPlot for true B0Bar', 100, -1.0,
                                  1.0)
    # belle plot with true B0 and B0bars
    hallo12 = ROOT.TH1F('BellePlot_NoCut', 'BellePlot_NoCut (binning 100)',
                        100, -1.0, 1.0)

    ###############################
    diag = ROOT.TF1('diag', 'pol1', -1, 1)
    ###############################

    # histograms for the efficiency calculation in wrong way
    histo_m0 = ROOT.TH1F('BellePlot_m0',
                         'BellePlot_m for true B0 (binning 50)', 50, -1.0, 1.0)
    histo_m1 = ROOT.TH1F('BellePlot_m1',
                         'BellePlot_m for true B0 (binning 50)', 50, -1.0, 1.0)
    histo_m2 = ROOT.TH1F('BellePlot_m2',
                         'BellePlot_m for true B0Bar (binning 50)', 50, -1.0,
                         1.0)

    # filling the histograms

    tree.Draw(method + '_qrCombined>>BellePlot_B0', 'qrMC == 1')
    tree.Draw(method + '_qrCombined>>BellePlot_B0Bar', 'qrMC == -1')
    tree.Draw(method + '_qrCombined>>BellePlot_NoCut', 'abs(qrMC) == 1')

    tree.Draw(method + '_qrCombined>>Calibration_B0', 'qrMC == 1')
    tree.Draw(method + '_qrCombined>>Calibration_B0Bar', 'qrMC == -1')

    tree.Draw(method + '_qrCombined>>notTaggedEvents',
              'abs(qrMC) == 0 && isSignal == 1 && ' +
              method + '_qrCombined < -1')

    # filling histograms wrong efficiency calculation
    tree.Draw(method + '_qrCombined>>BellePlot_m0',
              'qrMC == 1 && ' + method + '_qrCombined>0')
    tree.Draw(method + '_qrCombined>>BellePlot_m1',
              'qrMC == 1 && ' + method + '_qrCombined<0')
    tree.Draw(method + '_qrCombined>>BellePlot_m2',
              'qrMC == -1 && ' + method + '_qrCombined>0 ')

    # filling with abs(qr) in one of 6 bins with its weight
    # separate calculation for B0 and B0bar

    tree.Project('Average_r', 'abs(' + method + '_qrCombined)', 'abs(' + method + '_qrCombined)*(abs(qrMC) == 1)')
    tree.Project('Average_rB0', 'abs(' + method + '_qrCombined)', 'abs(' + method + '_qrCombined)*(qrMC==1)')
    tree.Project('Average_rB0bar', 'abs(' + method + '_qrCombined)', 'abs(' + method + '_qrCombined)*(qrMC==-1)')

    tree.Project('MS_r', 'abs(' + method + '_qrCombined)', '(' + method +
                 '_qrCombined*' + method + '_qrCombined)*(abs(qrMC) == 1)')
    tree.Project('MS_rB0', 'abs(' + method + '_qrCombined)',
                 '(' + method + '_qrCombined*' + method + '_qrCombined)*(qrMC==1)')
    tree.Project('MS_rB0bar', 'abs(' + method + '_qrCombined)',
                 '(' + method + '_qrCombined*' + method + '_qrCombined)*(qrMC==-1)')

    # filling with abs(qr) in one of 6 bins
    tree.Project('entries_per_bin', 'abs(' + method + '_qrCombined)', 'abs(qrMC) == 1')
    tree.Project('entries_per_binB0', 'abs(' + method + '_qrCombined)', 'qrMC == 1')
    tree.Project('entries_per_binB0bar', 'abs(' + method + '_qrCombined)', 'qrMC == -1')

    # filling histograms with number of wrong tags per r-bin
    tree.Project('mc_NwB0', 'abs(' + method + '_qrCombined)', ' ' + method + '_qrCombined*qrMC < 0 && qrMC == 1')
    tree.Project('mc_NwB0bar', 'abs(' + method + '_qrCombined)', ' ' + method + '_qrCombined*qrMC < 0 && qrMC == -1')

    # producing the average r histograms
    histo_avr_r.Divide(histo_entries_per_bin)
    histo_avr_rB0.Divide(histo_entries_per_binB0)
    histo_avr_rB0bar.Divide(histo_entries_per_binB0bar)

    histo_ms_r.Divide(histo_entries_per_bin)
    histo_ms_rB0.Divide(histo_entries_per_binB0)
    histo_ms_rB0bar.Divide(histo_entries_per_binB0bar)

    # producing the calibration plots
    # Errors ok
    histo_calib_B0.Divide(hallo12)
    histo_calib_B0bar.Divide(hallo12)

    # Fit for calibration plot
    print(' ')
    print('****************************** CALIBRATION CHECK FOR COMBINER USING ' +
          method + ' ***************************************')
    print(' ')
    print('Fit ploynomial of first order to the calibration plot. Expected value ~0.5')
    print(' ')
    histo_calib_B0.Fit(diag, 'TEST')
    print('       ')
    print('****************************** MEASURED EFFECTIVE EFFICIENCY FOR COMBINER USING ' +
          method + ' ***********************************')
    print('*                                                                               ' +
          '                                       *')
    # get total number of entries
    total_tagged = histo_entries_per_bin.GetEntries()
    total_tagged_B0 = histo_entries_per_binB0.GetEntries()
    total_tagged_B0bar = histo_entries_per_binB0bar.GetEntries()
    total_notTagged = histo_notTaggedEvents.GetEntries()
    total_entries = (total_tagged + total_notTagged)
    # To a good approximation we assume that half of the not tagged B mesons were B0 (B0bar)
    total_entriesB0 = (total_tagged_B0 + total_notTagged / 2)
    total_entriesB0bar = (total_tagged_B0bar + total_notTagged / 2)

    tagging_eff = total_tagged / (total_tagged + total_notTagged)
    DeltaTagging_eff = math.sqrt(total_tagged * total_notTagged**2 + total_notTagged * total_tagged**2) / (total_entries**2)
    tot_eff_effB0 = 0
    tot_eff_effB0bar = 0
    average_eff_eff = 0
    uncertainty_eff_effB0 = 0
    uncertainty_eff_effB0bar = 0
    uncertainty_eff_effAverage = 0
    diff_eff_Uncertainty = 0
    event_fractionB0 = array('f', [0] * r_size)
    event_fractionB0bar = array('f', [0] * r_size)
    event_fractionTotal = array('f', [0] * r_size)
    event_fractionTotalUncertainty = array('f', [0] * r_size)
    eventsInBin_B0 = array('f', [0] * r_size)
    eventsInBin_B0bar = array('f', [0] * r_size)
    eventsInBin_Total = array('f', [0] * r_size)
    event_fractionDiff = array('f', [0] * r_size)
    event_fractionDiffUncertainty = array('f', [0] * r_size)
    rvalueB0 = array('f', [0] * r_size)
    rvalueB0bar = array('f', [0] * r_size)
    rvalueB0Average = array('f', [0] * r_size)
    rvalueStdB0 = array('f', [0] * r_size)
    rvalueStdB0bar = array('f', [0] * r_size)
    rvalueStdB0Average = array('f', [0] * r_size)
    wvalue = array('f', [0] * r_size)
    wvalueUncertainty = array('f', [0] * r_size)
    wvalueB0 = array('f', [0] * r_size)
    wvalueB0bar = array('f', [0] * r_size)
    wvalueB0Uncertainty = array('f', [0] * r_size)
    wvalueB0barUncertainty = array('f', [0] * r_size)
    wvalueDiff = array('f', [0] * r_size)
    wvalueDiffUncertainty = array('f', [0] * r_size)
    entries = array('f', [0] * r_size)
    entriesB0 = array('f', [0] * r_size)
    entriesB0bar = array('f', [0] * r_size)
    iEffEfficiency = array('f', [0] * r_size)
    iEffEfficiencyUncertainty = array('f', [0] * r_size)
    iEffEfficiencyB0Uncertainty = array('f', [0] * r_size)
    iEffEfficiencyB0barUncertainty = array('f', [0] * r_size)
    iEffEfficiencyB0UncertaintyFromOutput = array('f', [0] * r_size)
    iEffEfficiencyB0barUncertaintyFromOutput = array('f', [0] * r_size)

    iDeltaEffEfficiency = array('f', [0] * r_size)
    iDeltaEffEfficiencyUncertainty = array('f', [0] * r_size)
    muParam = array('f', [0] * r_size)
    muParamUncertainty = array('f', [0] * r_size)
    # intervallEff = array('f', [0] * r_size)

    print('*                 -->  DETERMINATION BASED ON MONTE CARLO           ' +
          '                                                   *')
    print('*                                                                   ' +
          '                                                   *')
    print('*  Note: mu = Delta_Effcy/(2*Efficiency). Needed for CP analysis ' +
          'together with w and  Delta_w                          *')
    print('*                                                                   ' +
          '                                                   *')
    print('* ------------------------------------------------------------------' +
          '-------------------------------------------------- *')
    print('*   r-interval          <r>        Efficiency   Delta_Effcy      ' +
          '     mu                 w               Delta_w       *')
    print('* ------------------------------------------------------------------' +
          '-------------------------------------------------- *')
    performance = []
    for i in range(1, r_size):
        # get the average r-value
        entries[i] = histo_entries_per_bin.GetBinContent(i)
        entriesB0[i] = histo_entries_per_binB0.GetBinContent(i)
        entriesB0bar[i] = histo_entries_per_binB0bar.GetBinContent(i)
        # fraction of events/all events

        event_fractionB0[i] = entriesB0[i] / total_entriesB0
        event_fractionB0bar[i] = entriesB0bar[i] / total_entriesB0bar

        # event_fractionTotal[i] = (entriesB0[i] + entriesB0bar[i]) / total_entries
        # event_fractionDiff[i] = (entriesB0[i] - entriesB0bar[i]) / total_entries

        event_fractionTotal[i] = (event_fractionB0[i] + event_fractionB0bar[i]) / 2
        event_fractionDiff[i] = event_fractionB0[i] - event_fractionB0bar[i]

        event_fractionDiffUncertainty[i] = math.sqrt(entriesB0[i] *
                                                     (total_entriesB0 -
                                                      entriesB0[i]) /
                                                     total_entriesB0**3 +
                                                     entriesB0bar[i] *
                                                     (total_entriesB0bar -
                                                      entriesB0bar[i]) /
                                                     total_entriesB0bar**3)

        event_fractionTotalUncertainty[i] = event_fractionDiffUncertainty[i] / 2

        rvalueB0[i] = histo_avr_rB0.GetBinContent(i)
        rvalueB0bar[i] = histo_avr_rB0bar.GetBinContent(i)
        rvalueB0Average[i] = histo_avr_r.GetBinContent(i)  # (rvalueB0[i] + rvalueB0bar[i]) / 2
        rvalueStdB0[i] = math.sqrt(histo_ms_rB0.GetBinContent(
            i) - (histo_avr_rB0.GetBinContent(i))**2) / math.sqrt(entriesB0[i] - 1)
        rvalueStdB0bar[i] = math.sqrt(histo_ms_rB0bar.GetBinContent(
            i) - (histo_avr_rB0bar.GetBinContent(i))**2) / math.sqrt(entriesB0bar[i] - 1)
        rvalueStdB0Average[i] = math.sqrt(rvalueStdB0[i]**2 + rvalueStdB0bar[i]**2) / 2
        # math.sqrt(histo_ms_r.GetBinContent(i) - (histo_avr_r.GetBinContent(i))**2)
        # calculate the wrong tag fractin (only true if MC data good)

        wvalueB0[i] = histo_mc_NwB0.GetBinContent(i) / entriesB0[i]  # (1 - rvalueB0[i]) / 2
        wvalueB0bar[i] = histo_mc_NwB0bar.GetBinContent(i) / entriesB0bar[i]  # (1 - rvalueB0bar[i]) / 2
        wvalueDiff[i] = wvalueB0[i] - wvalueB0bar[i]
        wvalueB0Uncertainty[i] = math.sqrt(histo_mc_NwB0.GetBinContent(
            i) * (entriesB0[i] - histo_mc_NwB0.GetBinContent(i)) / (entriesB0[i]**3))
        wvalueB0barUncertainty[i] = math.sqrt(histo_mc_NwB0bar.GetBinContent(
            i) * (entriesB0bar[i] - histo_mc_NwB0bar.GetBinContent(i)) / (entriesB0bar[i]**3))
        # math.sqrt((rvalueStdB0[i] / 2)**2 + (rvalueStdB0bar[i] / 2)**2)
        wvalueDiffUncertainty[i] = math.sqrt(wvalueB0Uncertainty[i]**2 + wvalueB0barUncertainty[i]**2)
        wvalue[i] = (wvalueB0[i] + wvalueB0bar[i]) / 2
        wvalueUncertainty[i] = wvalueDiffUncertainty[i] / 2

        # Avr Efficiency

        # iEffEfficiency[i] = (event_fractionB0[i] * rvalueB0[i] * rvalueB0[i] +
        #                                   event_fractionB0bar[i] * rvalueB0bar[i] * rvalueB0bar[i]) / 2

        iEffEfficiency[i] = event_fractionTotal[i] * (1 - 2 * wvalue[i])**2

        iEffEfficiencyUncertainty[i] = (1 - 2 * wvalue[i]) * \
            math.sqrt((2 * event_fractionTotal[i] * 2 * wvalueUncertainty[i])**2 +
                      (1 - 2 * wvalue[i])**2 * event_fractionTotalUncertainty[i]**2)

        # iEffEfficiencyUncertainty[i] = rvalueB0Average[i] * \
        #     math.sqrt((2 * total_entries * entries[i] * rvalueStdB0Average[i])**2 +
        #               rvalueB0Average[i]**2 * entries[i] *
        #               (total_entries * (total_entries - entries[i]) +
        #                entries[i] * total_notTagged)) / (total_entries**2)

        # iEffEfficiencyB0UncertaintyFromOutput[i] = rvalueB0[i] * \
        #     math.sqrt((2 * total_entriesB0 * entriesB0[i] * rvalueStdB0[i])**2 +
        #               rvalueB0[i]**2 * entriesB0[i] *
        #               total_entriesB0 * (total_entriesB0 - entriesB0[i])) / (total_entriesB0**2)
        # iEffEfficiencyB0barUncertaintyFromOutput[i] = rvalueB0bar[i] * \
        #     math.sqrt((2 * total_entriesB0bar * entriesB0bar[i] * rvalueStdB0bar[i])**2 +
        #               rvalueB0bar[i]**2 * entriesB0bar[i] *
        #               total_entriesB0bar * (total_entriesB0bar - entriesB0bar[i])) / (total_entriesB0bar**2)

        # iEffEfficiencyUncertainty[i] =
        # math.sqrt(iEffEfficiencyB0UncertaintyFromOutput[i]**2 +
        # iEffEfficiencyB0barUncertaintyFromOutput[i]**2) / 2

        # iEffEfficiency[i] = (event_fractionB0[i] * (1-2*wvalueB0[i])**2 + event_fractionB0bar[i] * (1-2*wvalueB0bar[i])**2)/2

        average_eff_eff += iEffEfficiency[i]
        # average_eff_eff += (event_fractionB0[i] * (1-2*wvalueB0[i])**2 + event_fractionB0bar[i] * (1-2*wvalueB0bar[i])**2)/2

        # Delta Eff

        iDeltaEffEfficiency[i] = event_fractionB0[i] * (1 - 2 * wvalueB0[i])**2 - \
            event_fractionB0bar[i] * (1 - 2 * wvalueB0bar[i])**2

        iEffEfficiencyB0Uncertainty[i] = (1 - 2 * wvalueB0[i]) * \
            math.sqrt((2 * total_entriesB0 * entriesB0[i] * 2 * wvalueB0Uncertainty[i])**2 +
                      (1 - 2 * wvalueB0[i])**2 * entriesB0[i] *
                      total_entriesB0 * (total_entriesB0 - entriesB0[i])) / (total_entriesB0**2)
        iEffEfficiencyB0barUncertainty[i] = (1 - 2 * wvalueB0bar[i]) * \
            math.sqrt((2 * total_entriesB0bar * entriesB0bar[i] * 2 * wvalueB0barUncertainty[i])**2 +
                      (1 - 2 * wvalueB0bar[i])**2 * entriesB0bar[i] *
                      total_entriesB0bar * (total_entriesB0bar - entriesB0bar[i])) / (total_entriesB0bar**2)

        iDeltaEffEfficiencyUncertainty[i] = math.sqrt(iEffEfficiencyB0Uncertainty[i]**2 + iEffEfficiencyB0barUncertainty[i]**2)
        # iEffEfficiencyUncertainty[i] = iDeltaEffEfficiencyUncertainty[i]/2

        diff_eff_Uncertainty = diff_eff_Uncertainty + iDeltaEffEfficiencyUncertainty[i]**2

        # finally calculating the total effective efficiency
        tot_eff_effB0 = tot_eff_effB0 + event_fractionB0[i] * (1 - 2 * wvalueB0[i])**2
        tot_eff_effB0bar = tot_eff_effB0bar + event_fractionB0bar[i] * (1 - 2 * wvalueB0bar[i])**2
        uncertainty_eff_effAverage = uncertainty_eff_effAverage + iEffEfficiencyUncertainty[i]**2
        uncertainty_eff_effB0 = uncertainty_eff_effB0 + iEffEfficiencyB0Uncertainty[i]**2
        uncertainty_eff_effB0bar = uncertainty_eff_effB0bar + iEffEfficiencyB0barUncertainty[i]**2
        muParam[i] = event_fractionDiff[i] / (2 * event_fractionTotal[i])
        muParamUncertainty[i] = event_fractionDiffUncertainty[i] / (2 * event_fractionTotal[i]) * math.sqrt(muParam[i]**2 + 1)

        # intervallEff[i] = event_fractionTotal[i] * rvalueB0Average[i] * rvalueB0Average[i]
        print('* ' + '{:.3f}'.format(r_subsample[i - 1]) + ' - ' + '{:.3f}'.format(r_subsample[i]) + '   ' +
              '{:.3f}'.format(rvalueB0Average[i]) + ' +- ' + '{:.4f}'.format(rvalueStdB0Average[i]) + '    ' +
              '{:.4f}'.format(event_fractionTotal[i]) + ' ' +
              '{: .4f}'.format(event_fractionDiff[i]) + ' +- ' + '{:.4f}'.format(event_fractionDiffUncertainty[i]) + ' ' +
              '{: .4f}'.format(muParam[i]) + ' +- ' + '{:.4f}'.format(muParamUncertainty[i]) + '   ' +
              '{:.4f}'.format(wvalue[i]) + ' +- ' + '{:.4f}'.format(wvalueUncertainty[i]) + '  ' +
              '{: .4f}'.format(wvalueDiff[i]) + ' +- ' + '{:.4f}'.format(wvalueDiffUncertainty[i]) + '  *')

    # average_eff_eff = (tot_eff_effB0 + tot_eff_effB0bar) / 2
    uncertainty_eff_effAverage = math.sqrt(uncertainty_eff_effAverage)
    uncertainty_eff_effB0 = math.sqrt(uncertainty_eff_effB0)
    uncertainty_eff_effB0bar = math.sqrt(uncertainty_eff_effB0bar)
    diff_eff = tot_eff_effB0 - tot_eff_effB0bar
    diff_eff_Uncertainty = math.sqrt(diff_eff_Uncertainty)
    print('* --------------------------------------------------------------------------------------------------' +
          '------------------ *')
    print('*                                                                                *')
    print('*    __________________________________________________________________________________________    *')
    print('*   |                                                                                |   *')
    print('*   | TOTAL NUMBER OF TAGGED EVENTS  =  ' +
          '{:<24}'.format("%.0f" % total_tagged) + '{:>36}'.format('|   *'))
    print('*   |                                                                                |   *')
    print(
        '*   | TOTAL AVERAGE EFFICIENCY  (q=+-1)=  ' +
        '{:.2f}'.format(
            tagging_eff *
            100) +
        " +- " +
        '{:.2f}'.format(
            DeltaTagging_eff *
            100) +
        ' %                                      |   *')
    print('*   |                                                                                |   *')
    print(
        '*   | TOTAL AVERAGE EFFECTIVE EFFICIENCY  (q=+-1)=  ' +
        '{:.6f}'.format(
            average_eff_eff *
            100) +
        " +- " +
        '{:.6f}'.format(
            uncertainty_eff_effAverage *
            100) +
        ' %                    |   *')
    print('*   |                                                                                |   *')
    print(
        '*   | TOTAL AVERAGE EFFECTIVE EFFICIENCY ASYMMETRY (q=+-1)=  ' +
        '{:^9.6f}'.format(
            diff_eff *
            100) +
        " +- " +
        '{:.6f}'.format(
            diff_eff_Uncertainty *
            100) +
        ' %           |   *')
    print('*   |                                                                                |   *')
    print('*   | B0-TAGGER  TOTAL EFFECTIVE EFFICIENCIES: ' +
          '{:.2f}'.format(tot_eff_effB0 * 100) + " +-" + '{: 4.2f}'.format(uncertainty_eff_effB0 * 100) +
          ' % (q=+1)  ' +
          '{:.2f}'.format(tot_eff_effB0bar * 100) + " +-" + '{: 4.2f}'.format(uncertainty_eff_effB0bar * 100) +
          ' % (q=-1) ' + ' |   *')
    print('*   |                                                                                |   *')
    print('*   | FLAVOR PERCENTAGE (MC):                  ' +
          '{:.2f}'.format(total_tagged_B0 / total_tagged * 100) + ' % (q=+1)  ' +
          '{:.2f}'.format(total_tagged_B0bar / total_tagged * 100) + ' % (q=-1)  Diff=' +
          '{:^5.2f}'.format((total_tagged_B0 - total_tagged_B0bar) / total_tagged * 100) + ' %    |   *')
    print('*   |__________________________________________________________________________________________|   *')
    print('*                                                                                *')
    print('****************************************************************************************************')
    print('*                                                                                *')

    # not that imortant
    print('* ---------------------------------                                                                *')
    print('* Efficiency Determination - easiest way                                                           *')
    print('* ---------------------------------                                                                *')
    total_tagged_B0 = histo_belleplotB0.GetEntries()
    total_tagged_B0Bar = histo_belleplotB0bar.GetEntries()
    total_tagged_wrong = histo_m1.GetEntries()
    total_tagged_B0Bar_wrong = histo_m2.GetEntries()
    total_tagged = total_tagged_B0 + total_tagged_B0Bar
    total_tagged_wrong = total_tagged_wrong + total_tagged_B0Bar_wrong

    wrong_tag_fraction_B0 = total_tagged_wrong / total_tagged_B0
    wrong_tag_fraction_B0Bar = total_tagged_B0Bar_wrong / total_tagged_B0Bar
    wrong_tag_fraction = total_tagged_wrong / total_tagged
    right_tag_fraction_B0 = 1 - 2 * wrong_tag_fraction_B0
    right_tag_fraction_B0Bar = 1 - 2 * wrong_tag_fraction_B0Bar
    right_tag_fraction = 1 - 2 * wrong_tag_fraction
    wrong_eff_B0 = right_tag_fraction_B0 * right_tag_fraction_B0
    wrong_eff_B0Bar = right_tag_fraction_B0Bar * right_tag_fraction_B0Bar
    wrong_eff = right_tag_fraction * right_tag_fraction

    print('*     wrong_tag_fraction for all:   ' +
          '{:.3f}'.format(wrong_tag_fraction * 100) +
          ' %                                                       *')
    print('*     right_tag_fraction for all:   ' +
          '{:.3f}'.format(right_tag_fraction * 100) +
          ' %                                                       *')
    print('*     wrong calculated eff all:     ' + '{:.3f}'.format(wrong_eff * 100) +
          ' %                                                       *')
    print('*                                                                                *')
    print('****************************************************************************************************')
    print('')
    print('Table For B2TIP')
    print('')
    # write out the histograms
    # histo_avr_r.Write('', ROOT.TObject.kOverwrite)
    # histo_entries_per_bin.Write('', ROOT.TObject.kOverwrite)

    # histo_Cnet_output_B0.Write('', ROOT.TObject.kOverwrite)
    # histo_Cnet_output_B0bar.Write('', ROOT.TObject.kOverwrite)
    # histo_belleplotB0.Write('', ROOT.TObject.kOverwrite)
    # histo_belleplotB0bar.Write('', ROOT.TObject.kOverwrite)
    # histo_calib_B0.Write('', ROOT.TObject.kOverwrite)
    # histo_calib_B0bar.Write('', ROOT.TObject.kOverwrite)

    maxB0 = histo_belleplotB0.GetBinContent(histo_belleplotB0.GetMaximumBin())
    maxB0bar = histo_belleplotB0bar.GetBinContent(histo_belleplotB0bar.GetMaximumBin())

    Ymax = max(maxB0, maxB0bar)
    Ymax = Ymax + Ymax / 12

    if YmaxForQrPlot < Ymax:
        YmaxForQrPlot = Ymax

    # produce a pdf
    ROOT.gStyle.SetOptStat(0)
    Canvas1 = ROOT.TCanvas('Bla' + method, 'Final Output', 1200, 800)
    Canvas1.cd()  # activate
    histo_belleplotB0.SetFillColorAlpha(ROOT.kBlue, 0.2)
    histo_belleplotB0.SetFillStyle(1001)
    histo_belleplotB0.GetYaxis().SetLabelSize(0.03)
    histo_belleplotB0.GetYaxis().SetLimits(0, YmaxForQrPlot)
    histo_belleplotB0.GetYaxis().SetTitleOffset(1.2)
    histo_belleplotB0.SetLineColor(ROOT.kBlue)
    histo_belleplotB0bar.SetFillColorAlpha(ROOT.kRed, 1.0)
    histo_belleplotB0bar.SetFillStyle(3005)
    histo_belleplotB0bar.SetLineColor(ROOT.kRed)
    # SetLabelSize etc SetTitle

    histo_belleplotB0.SetTitle('Final Flavor Tagger Output; #it{qr}-output ; Events'
                               )
    histo_belleplotB0.SetMinimum(0)
    histo_belleplotB0.SetMaximum(YmaxForQrPlot)
    histo_belleplotB0.Draw('hist')
    histo_belleplotB0bar.Draw('hist same')

    leg = ROOT.TLegend(0.75, 0.8, 0.9, 0.9)
    leg.AddEntry(histo_belleplotB0, 'true B0')
    leg.AddEntry(histo_belleplotB0bar, 'true B0bar')
    leg.Draw()

    Canvas1.Update()
    # IPython.embed()
    with Quiet(ROOT.kError):
        Canvas1.SaveAs(workingDirectory + '/' + 'PIC_Belleplot_both' + method + '.pdf')

    # produce the nice calibration plot
    Canvas2 = ROOT.TCanvas('Bla2' + method, 'Calibration plot for true B0', 1200, 800)
    Canvas2.cd()  # activate
    histo_calib_B0.SetFillColorAlpha(ROOT.kBlue, 0.2)
    histo_calib_B0.SetFillStyle(1001)
    histo_calib_B0.GetYaxis().SetTitleOffset(1.2)
    histo_calib_B0.SetLineColor(ROOT.kBlue)

    histo_calib_B0.SetTitle('Calibration For True B0; #it{qr}-output ; Calibration '
                            )
    histo_calib_B0.Draw('hist')
    diag.Draw('SAME')
    Canvas2.Update()
    with Quiet(ROOT.kError):
        Canvas2.SaveAs(workingDirectory + '/' + 'PIC_Calibration_B0' + method + '.pdf')

    histo_avr_r.Delete()
    histo_avr_rB0.Delete()
    histo_avr_rB0bar.Delete()
    histo_ms_r.Delete()
    histo_ms_rB0.Delete()
    histo_ms_rB0bar.Delete()
    histo_mc_NwB0.Delete()
    histo_mc_NwB0bar.Delete()
    histo_notTaggedEvents.Delete()
    histo_entries_per_bin.Delete()
    histo_entries_per_binB0.Delete()
    histo_entries_per_binB0bar.Delete()
    histo_Cnet_output_B0.Delete()
    histo_Cnet_output_B0bar.Delete()
    histo_belleplotB0.Delete()
    histo_belleplotB0bar.Delete()
    histo_calib_B0.Delete()
    histo_calib_B0bar.Delete()
    hallo12.Delete()
    histo_m0.Delete()
    histo_m1.Delete()
    histo_m2.Delete()
    Canvas1.Clear()
    Canvas2.Clear()

    print(r'\begin{tabularx}{1\textwidth}{@{}r  r  r  r  r  r  r@{}}')
    print(r'\hline')
    print(r'$r$- Interval $\enskip$ & $\varepsilon_i\ $ &  $\Delta\varepsilon_i\ $ & $w_i \pm \delta w_i\enskip\; $ ' +
          r' & $\Delta w_i \pm \delta\Delta w_i $& $\varepsilon_{\text{eff}, i} \pm \delta\varepsilon_{\text{eff}, i}\enskip\, $ ' +
          r' & $\Delta \varepsilon_{\text{eff}, i}  \pm \delta\Delta \varepsilon_{\text{eff}, i}\enskip\, $\\ \hline\hline')
    for i in range(1, r_size):
        print('$ ' + '{:.3f}'.format(r_subsample[i - 1]) + ' - ' + '{:.3f}'.format(r_subsample[i]) + '$ & $'
              '{: 6.1f}'.format(event_fractionTotal[i] * 100) + r'$ & $' +
              '{: 7.2f}'.format(event_fractionDiff[i] * 100) + r'\;$ & $' +
              '{: 7.2f}'.format(wvalue[i] * 100) + r" \pm " + '{:2.2f}'.format(wvalueUncertainty[i] * 100) + r'\enskip $ & $' +
              '{: 7.2f}'.format(wvalueDiff[i] * 100) + r" \pm " +
              '{:2.2f}'.format(wvalueDiffUncertainty[i] * 100) + r'\enskip $ & $' +
              '{: 8.4f}'.format(iEffEfficiency[i] * 100) +  # + '$ & $' +
              r" \pm " + '{:2.4f}'.format(iEffEfficiencyUncertainty[i] * 100) + r'\, $ & $' +
              '{: 6.4f}'.format(iDeltaEffEfficiency[i] * 100) +  # +
              r" \pm " + '{:2.4f}'.format(iDeltaEffEfficiencyUncertainty[i] * 100) +
              r'\enskip\enskip $ \\ ')
    print(r'\hline\hline')
    print(r'\multicolumn{1}{r}{Total} &  & \multicolumn{5}{r}{ $\varepsilon_\text{eff} = ' +
          r'\sum_i \varepsilon_i \cdot \langle 1-2w_i\rangle^2 = ' +
          '{: 6.2f}'.format(average_eff_eff * 100) + r" \pm " + '{: 6.2f}'.format(uncertainty_eff_effAverage * 100) + r'\enskip\, ')
    print(r'\Delta \varepsilon_\text{eff} = ' +
          '{: 6.2f}'.format(diff_eff * 100) + r" \pm " + '{: 6.2f}'.format(diff_eff_Uncertainty * 100) + r'\quad\  $ }' +
          r' \\ ')
    print(r'\hline')
    print(r'\end{tabular}')
    print('')

    print('')
    print('Mu-Values for Table')
    print('')

    for i in range(1, r_size):
        print('$ ' + '{:.3f}'.format(r_subsample[i - 1]) + ' - ' + '{:.3f}'.format(r_subsample[i]) + '$ & $'
              '{: 7.2f}'.format(muParam[i] * 100) + r" \pm " + '{:2.2f}'.format(muParamUncertainty[i] * 100) + r' $ & ')
# ************************************************
# DETERMINATION OF INDIVIDUAL EFFECTIVE EFFICIENCY
# ************************************************

# keep in mind:
# the individual efficiency is determined on basis of the combiner training.
# Whereas the efficiency is determined on basis of the final expert output.

print(ft.getEventLevelParticleLists(usedCategories))

print('******************************************* MEASURED EFFECTIVE EFFICIENCY FOR INDIVIDUAL CATEGORIES ' +
      '**********************************************')
print('*                                                                                ' +
      '                                             *')
# input: Classifier input from event-level. Output of event-level is recalculated for input on combiner-level.
# but is re-evaluated under combiner target. Signal is B0, background is B0Bar.
categoriesPerformance = []
NbinsCategories = 100
for category in usedCategories:
    # histogram of input variable (only signal) - not yet a probability! It's a classifier plot!
    hist_both = ROOT.TH1F('Both_' + category, 'Input Both (B0) ' +
                          category + ' (binning)', NbinsCategories, -1.0, 1.0)
    # histogram of input variable (only signal) - not yet a probability! It's a classifier plot!
    hist_signal = ROOT.TH1F('Signal_' + category, 'Input Signal (B0) ' +
                            category + ' (binning)', NbinsCategories, -1.0, 1.0)
    # histogram of input variable (only background) - not yet a probability! It's a classifier plot!
    hist_background = ROOT.TH1F('Background_' + category, 'Input Background (B0bar) ' +
                                category + ' (binning)', NbinsCategories, -1.0, 1.0)

    # per definition that input is not comparable to the network output, this has to be transformed.
    # probability output from 0 to 1 (corresponds to net output probability) -> calculation below
    hist_probB0 = ROOT.TH1F('Probability' + category,
                            'Transformed to probability (B0) (' + category + ')',
                            NbinsCategories, 0.0, 1.0)
    hist_probB0bar = ROOT.TH1F('ProbabilityB0bar_' + category,
                               'Transformed to probability (B0bar) (' + category + ')',
                               NbinsCategories, 0.0, 1.0)
    # qp output from -1 to 1 -> transformation below
    hist_qrB0 = ROOT.TH1F('QR' + category, 'Transformed to qp (B0)(' +
                          category + ')', NbinsCategories, -1.0, 1.0)
    hist_qrB0bar = ROOT.TH1F('QRB0bar_' + category, 'Transformed to qp (B0bar) (' +
                             category + ')', NbinsCategories, -1.0, 1.0)
    # histogram for abs(qp), i.e. this histogram contains the r-values -> transformation below
    # also used to get the number of entries, sorted into 6 bins
    histo_entries_per_bin = ROOT.TH1F('entries_per_bin_' + category, 'Abs(qp)(B0) (' + category + ')', int(r_size - 2), r_subsample)
    histo_entries_per_binB0 = ROOT.TH1F('entries_per_bin' + category, 'Abs(qp)(B0) (' +
                                        category + ')', int(r_size - 2), r_subsample)
    histo_entries_per_binB0bar = ROOT.TH1F('entries_per_binB0bar_' + category,
                                           'Abs(qp) (B0bar) (' + category + ')', int(r_size - 2), r_subsample)

    # histogram contains at the end the average r values -> calculation below
    # sorted into r bins
    hist_avr_rB0 = ROOT.TH1F('Average_r' + category, 'Average r for B0' +
                             category, int(r_size - 2), r_subsample)
    hist_avr_rB0bar = ROOT.TH1F('Average_rB0bar_' + category, 'Average r for B0bar' +
                                category, int(r_size - 2), r_subsample)

    hist_ms_rB0 = ROOT.TH1F('AverageSqrdR' + category, 'Average r sqrd for B0' +
                            category, int(r_size - 2), r_subsample)
    hist_ms_rB0bar = ROOT.TH1F('AverageSqrdRB0bar_' + category, 'Average r sqrd for B0bar' +
                               category, int(r_size - 2), r_subsample)

    # ****** TEST OF CALIBRATION ******
    # for calibration plot we want to have
    hist_all = ROOT.TH1F('All_' + category, 'Input Signal (B0) and Background (B0Bar)' +
                         category + ' (binning 50)', 50, 0.0, 1.0)
    hist_calib_B0 = ROOT.TH1F('Calib_' + category, 'Calibration Plot for true B0' +
                              category + ' (binning 50)', 50, 0.0, 1.0)

    # fill both
    if category != "KaonNotWeighted" and category != "LambdaNotWeighted":
        #  if category == 'Lambda' or combinerVariable.find('weighted') == -1:
        tree.Draw('qp' + category + '>>Both_' + category, 'abs(qrMC) == 1.0')
        # fill signal
        tree.Draw('qp' + category + '>>Signal_' + category, 'qrMC == 1.0')
        # fill background
        tree.Draw('qp' + category + '>>Background_' + category, 'qrMC == -1.0')
        tree.Draw('qp' + category + '>>All_' + category, 'qrMC!=0')
        tree.Draw('qp' + category + '>>Calib_' + category, 'qrMC == 1.0')

        # elif combinerVariable.find('weighted') != -1:
        #     tree.Draw('extraInfo__boWeightedQpOf' + category + '__bc' + '>>Both_' + category, 'abs(qrMC) == 1.0')
        #     tree.Draw('extraInfo__boWeightedQpOf' + category + '__bc' + '>>Signal_' + category, 'qrMC == 1.0')
        #     tree.Draw('extraInfo__boWeightedQpOf' + category + '__bc' + '>>Background_' + category, 'qrMC == -1.0')
        #     tree.Draw('extraInfo__boWeightedQpOf' + category + '__bc' + '>>All_' + category, 'qrMC!=0')
        #     tree.Draw('extraInfo__boWeightedQpOf' + category + '__bc' + '>>Calib_' + category, 'qrMC == 1.0')
        #     category = category + 'W'

    elif category == "KaonNotWeighted":
        tree.Draw('extraInfo__boQpOfKaon__bc' + '>>Both_' + category, 'abs(qrMC) == 1.0')
        tree.Draw('extraInfo__boQpOfKaon__bc' + '>>Signal_' + category, 'qrMC == 1.0')
        tree.Draw('extraInfo__boQpOfKaon__bc' + '>>Background_' + category, 'qrMC == -1.0')
        tree.Draw('extraInfo__boQpOfKaon__bc' + '>>All_' + category, 'qrMC!=0')
        tree.Draw('extraInfo__boQpOfKaon__bc' + '>>Calib_' + category, 'qrMC == 1.0')

    elif category == "LambdaNotWeighted":
        tree.Draw('extraInfo__boQpOfLambda__bc' + '>>Both_' + category, 'abs(qrMC) == 1.0')
        tree.Draw('extraInfo__boQpOfLambda__bc' + '>>Signal_' + category, 'qrMC == 1.0')
        tree.Draw('extraInfo__boQpOfLambda__bc' + '>>Background_' + category, 'qrMC == -1.0')
        tree.Draw('extraInfo__boQpOfLambda__bc' + '>>All_' + category, 'qrMC!=0')
        tree.Draw('extraInfo__boQpOfLambda__bc' + '>>Calib_' + category, 'qrMC == 1.0')

    hist_calib_B0.Divide(hist_all)
    # ****** produce the input plots from combiner level ******

    maxSignal = hist_signal.GetBinContent(hist_signal.GetMaximumBin())
    maxBackground = hist_background.GetBinContent(hist_background.GetMaximumBin())

    Ymax = max(maxSignal, maxBackground)
    Ymax = Ymax + Ymax / 12

    ROOT.gStyle.SetOptStat(0)
    with Quiet(ROOT.kError):
        Canvas = ROOT.TCanvas('Bla', 'TITEL BLA', 1200, 800)
    Canvas.cd()  # activate
    Canvas.SetLogy()
    hist_signal.SetFillColorAlpha(ROOT.kBlue, 0.2)
    hist_signal.SetFillStyle(1001)
    hist_signal.SetTitleSize(0.1)
    hist_signal.GetXaxis().SetLabelSize(0.04)
    hist_signal.GetYaxis().SetLabelSize(0.04)
    hist_signal.GetXaxis().SetTitleSize(0.05)
    hist_signal.GetYaxis().SetTitleSize(0.05)
    hist_signal.GetXaxis().SetTitleOffset(0.95)
    hist_signal.GetYaxis().SetTitleOffset(1.1)
    hist_signal.GetYaxis().SetLimits(0, Ymax)
    hist_signal.SetLineColor(ROOT.kBlue)
    hist_background.SetFillColorAlpha(ROOT.kRed, 1.0)
    hist_background.SetFillStyle(3005)
    hist_background.GetYaxis().SetLimits(0, Ymax)
    hist_background.SetLineColor(ROOT.kRed)

    hist_signal.SetTitle(category + ' category; #it{qp}-Output ; Events')
    # hist_signal.SetMinimum(0)
    hist_signal.SetMaximum(Ymax)
    # hist_background.SetMinimum(0)
    hist_background.SetMaximum(Ymax)

    hist_signal.Draw('hist')
    hist_background.Draw('hist same')

    if category == 'MaximumPstar':
        legend = ROOT.TLegend(0.4, 0.75, 0.6, 0.9)
    else:
        legend = ROOT.TLegend(0.6, 0.75, 0.8, 0.9)
    legend.AddEntry(hist_signal, 'true B0')
    legend.AddEntry(hist_background, 'true B0bar')
    legend.SetTextSize(0.05)
    legend.Draw()

    Canvas.Update()
    with Quiet(ROOT.kError):
        Canvas.SaveAs(workingDirectory + '/' + 'PIC_' + category + '_Input_Combiner.pdf')

    # ***** TEST OF CALIBRATION ******

    # initialize some arrays
    binCounter = int(NbinsCategories + 1)
    dilutionB02 = array('d', [0] * binCounter)
    dilutionB0bar2 = array('d', [0] * binCounter)
    purityB0 = array('d', [0] * binCounter)
    purityB0bar = array('d', [0] * binCounter)
    signal = array('d', [0] * binCounter)
    back = array('d', [0] * binCounter)
    weight = array('d', [0] * binCounter)

    for i in range(1, binCounter):
        # doing the transformation to probabilities
        signal[i] = hist_signal.GetBinContent(i)
        back[i] = hist_background.GetBinContent(i)
        weight[i] = signal[i] + back[i]

        # avoid dividing by zero
        if signal[i] + back[i] == 0:
            purityB0[i] = 0
            dilutionB02[i] = 0
            purityB0bar[i] = 0
            dilutionB0bar2[i] = 0
        else:

            purityB0[i] = signal[i] / (signal[i] + back[i])
            dilutionB02[i] = -1 + 2 * signal[i] / (signal[i] + back[i])

            purityB0bar[i] = back[i] / (signal[i] + back[i])
            dilutionB0bar2[i] = -1 + 2 * back[i] / (signal[i] + back[i])

        # filling histogram with probability from 0 to 1
        hist_probB0.Fill(purityB0[i], signal[i])
        hist_probB0bar.Fill(purityB0bar[i], back[i])

        # filling histogram with qr from -1 to 1
        hist_qrB0.Fill(dilutionB02[i], signal[i])
        hist_qrB0bar.Fill(dilutionB0bar2[i], back[i])

        # filling histogram with abs(qr), i.e. this histogram contains the r-values (not qp)
        histo_entries_per_binB0.Fill(abs(dilutionB02[i]), signal[i])
        histo_entries_per_binB0bar.Fill(abs(dilutionB0bar2[i]), back[i])
        # filling histogram with abs(qr) special weighted - needed for average r calculation
        hist_avr_rB0.Fill(abs(dilutionB02[i]), abs(dilutionB02[i]) * signal[i])
        hist_avr_rB0bar.Fill(abs(dilutionB0bar2[i]), abs(dilutionB0bar2[i]) * back[i])
        # filling histogram with abs(qr)**2 special weighted - needed for std dev of r calculation
        hist_ms_rB0.Fill(abs(dilutionB02[i]), abs(dilutionB02[i] * dilutionB02[i]) * signal[i])
        hist_ms_rB0bar.Fill(abs(dilutionB0bar2[i]), abs(dilutionB0bar2[i] * dilutionB02[i]) * back[i])

    # hist_avr_rB0bar contains now the average r-value
    hist_avr_rB0.Divide(histo_entries_per_binB0)
    hist_avr_rB0bar.Divide(histo_entries_per_binB0bar)

    hist_ms_rB0.Divide(histo_entries_per_binB0)
    hist_ms_rB0bar.Divide(histo_entries_per_binB0bar)
    # now calculating the efficiency

    # calculating number of events
    total_entriesB0 = total_notTagged / 2
    total_entriesB0bar = total_notTagged / 2
    for i in range(1, r_size):
        total_entriesB0 = total_entriesB0 + histo_entries_per_binB0.GetBinContent(i)
        total_entriesB0bar = total_entriesB0bar + histo_entries_per_binB0bar.GetBinContent(i)
    # initializing some arrays
    tot_eff_effB0 = 0
    tot_eff_effB0bar = 0
    uncertainty_eff_effB0 = 0
    uncertainty_eff_effB0bar = 0
    uncertainty_eff_effAverage = 0
    diff_eff_Uncertainty = 0
    event_fractionB0 = array('f', [0] * r_size)
    event_fractionB0bar = array('f', [0] * r_size)
    rvalueB0 = array('f', [0] * r_size)
    rvalueB0bar = array('f', [0] * r_size)
    rvalueStdB0 = array('f', [0] * r_size)
    rvalueStdB0bar = array('f', [0] * r_size)
    # wvalue = array('f', [0] * r_size)
    entriesBoth = array('f', [0] * r_size)
    entriesB0 = array('f', [0] * r_size)
    entriesB0bar = array('f', [0] * r_size)
    iEffEfficiencyB0Uncertainty = array('f', [0] * r_size)
    iEffEfficiencyB0barUncertainty = array('f', [0] * r_size)
    iDeltaEffEfficiencyUncertainty = array('f', [0] * r_size)

    for i in range(1, r_size):

        entriesBoth[i] = entriesB0bar[i] + entriesB0[i]
        entriesB0[i] = histo_entries_per_binB0.GetBinContent(i)
        entriesB0bar[i] = histo_entries_per_binB0bar.GetBinContent(i)
        event_fractionB0[i] = entriesB0[i] / total_entriesB0
        event_fractionB0bar[i] = entriesB0bar[i] / total_entriesB0bar
        # print '*  Bin ' + str(i) + ' r-value: ' + str(rvalueB0[i]), 'entriesB0: ' +
        # str(event_fractionB0[i] * 100) + ' % (' + str(entriesB0[i]) + '/' +
        # str(total_entriesB0) + ')'

        rvalueB0[i] = hist_avr_rB0.GetBinContent(i)
        rvalueB0bar[i] = hist_avr_rB0bar.GetBinContent(i)

        rvalueStdB0[i] = 0
        rvalueStdB0bar[i] = 0

        if entriesB0[i] > 1:
            rvalueStdB0[i] = math.sqrt(abs(hist_ms_rB0.GetBinContent(
                i) - (hist_avr_rB0.GetBinContent(i))**2)) / math.sqrt(entriesB0[i] - 1)

        if entriesB0bar[i] > 1:
            rvalueStdB0bar[i] = math.sqrt(abs(hist_ms_rB0bar.GetBinContent(
                i) - (hist_avr_rB0bar.GetBinContent(i))**2)) / math.sqrt(entriesB0bar[i] - 1)
        # wvalue[i] = (1 - rvalueB0[i]) / 2

        tot_eff_effB0 = tot_eff_effB0 + event_fractionB0[i] * rvalueB0[i] \
            * rvalueB0[i]
        tot_eff_effB0bar = tot_eff_effB0bar + event_fractionB0bar[i] * rvalueB0bar[i] \
            * rvalueB0bar[i]

        iEffEfficiencyB0Uncertainty[i] = rvalueB0[i] * \
            math.sqrt((2 * total_entriesB0 * entriesB0[i] * rvalueStdB0[i])**2 +
                      rvalueB0[i]**2 * entriesB0[i] *
                      (total_entriesB0 * (total_entriesB0 - entriesB0[i]) +
                       entriesB0[i] * total_notTagged)) / (total_entriesB0**2)
        iEffEfficiencyB0barUncertainty[i] = rvalueB0bar[i] * \
            math.sqrt((2 * total_entriesB0bar * entriesB0bar[i] * rvalueStdB0bar[i])**2 +
                      rvalueB0bar[i]**2 * entriesB0bar[i] *
                      (total_entriesB0bar * (total_entriesB0bar - entriesB0bar[i]) +
                       entriesB0bar[i] * total_notTagged)) / (total_entriesB0bar**2)

        iDeltaEffEfficiencyUncertainty[i] = math.sqrt(iEffEfficiencyB0Uncertainty[i]**2 + iEffEfficiencyB0barUncertainty[i]**2)

        diff_eff_Uncertainty = diff_eff_Uncertainty + iDeltaEffEfficiencyUncertainty[i]**2

        uncertainty_eff_effB0 = uncertainty_eff_effB0 + iEffEfficiencyB0Uncertainty[i]**2
        uncertainty_eff_effB0bar = uncertainty_eff_effB0bar + iEffEfficiencyB0barUncertainty[i]**2

    effDiff = tot_eff_effB0 - tot_eff_effB0bar
    effAverage = (tot_eff_effB0 + tot_eff_effB0bar) / 2

    uncertainty_eff_effB0 = math.sqrt(uncertainty_eff_effB0)
    uncertainty_eff_effB0bar = math.sqrt(uncertainty_eff_effB0bar)
    diff_eff_Uncertainty = math.sqrt(diff_eff_Uncertainty)
    uncertainty_eff_effAverage = diff_eff_Uncertainty / 2
    print(
        '{:<25}'.format("* " + category) + ' B0-Eff=' +
        '{: 8.2f}'.format(tot_eff_effB0 * 100) + " +-" + '{: 4.2f}'.format(uncertainty_eff_effB0 * 100) +
        ' %' +
        '   B0bar-Eff=' +
        '{: 8.2f}'.format(tot_eff_effB0bar * 100) + " +-" + '{: 4.2f}'.format(uncertainty_eff_effB0bar * 100) +
        ' %' +
        '   EffAverage=' +
        '{: 8.2f}'.format(effAverage * 100) + " +- " + '{:4.2f}'.format(uncertainty_eff_effAverage * 100) + ' %' +
        '   EffDiff=' +
        '{: 8.2f}'.format(effDiff * 100) + " +- " + '{:4.2f}'.format(diff_eff_Uncertainty * 100) + ' %  *')

    # hist_signal.Write('', ROOT.TObject.kOverwrite)
    # hist_background.Write('', ROOT.TObject.kOverwrite)
    # hist_probB0.Write('', ROOT.TObject.kOverwrite)
    # hist_probB0bar.Write('', ROOT.TObject.kOverwrite)
    # hist_qpB0.Write('', ROOT.TObject.kOverwrite)
    # hist_qpB0bar.Write('', ROOT.TObject.kOverwrite)
    # hist_absqpB0.Write('', ROOT.TObject.kOverwrite)
    # hist_absqpB0bar.Write('', ROOT.TObject.kOverwrite)
    # hist_avr_rB0.Write('', ROOT.TObject.kOverwrite)
    # hist_avr_rB0bar.Write('', ROOT.TObject.kOverwrite)
    # hist_all.Write('', ROOT.TObject.kOverwrite)
    # hist_calib_B0.Write('', ROOT.TObject.kOverwrite)
    categoriesPerformance.append((category, effAverage, uncertainty_eff_effAverage, effDiff, diff_eff_Uncertainty))
    with Quiet(ROOT.kError):
        Canvas.Clear()
# if average_eff != 0:
    # print '*    -------------------------------------------------------------------------'
    # print '*    ' + '{: > 8.2f}'.format(average_eff * 100) + ' %' \
    # + '{:>85}'.format('TOTAL' + '                      *')

print('*                                                                                ' +
      '                       *')
print('**************************************************************************************************************************' +
      '************************')
print('')
print('Table For B2TIP')
print('')
print(r'\begin{tabular}{ l  r  r }')
print(r'\hline')
print(r'Categories & $\varepsilon_\text{eff} \pm \delta\varepsilon_\text{eff} $& ' +
      r'$\Delta\varepsilon_\text{eff} \pm \delta\Delta\varepsilon_\text{eff}$\\ \hline\hline')
for (category, effAverage, uncertainty_eff_effAverage, effDiff, diff_eff_Uncertainty) in categoriesPerformance:
    print(
        '{:<23}'.format(category) +
        ' & $' +
        '{: 6.2f}'.format(effAverage * 100) + r" \pm " + '{:4.2f}'.format(uncertainty_eff_effAverage * 100) +
        ' $ & $' +
        '{: 6.2f}'.format(effDiff * 100) + r" \pm " + '{:4.2f}'.format(diff_eff_Uncertainty * 100) +
        r'\ \enskip  $ \\')
print(r'\hline')
print(r'\end{tabular}')
B2INFO('qp Output Histograms in pdf format saved at: ' + workingDirectory)
with open("categoriesPerformance.pkl", "wb") as f:
    pickle.dump(categoriesPerformance, f)


# 281
from itertools import *

def double(n):
	for i in range(1, len(n)):
		if n[i] == n[i - 1]:
			return False
	return True

def alternation(n):
	for i in range(1, len(n)):
		if int(n[i]) % 2 == int(n[i - 1]) % 2:
			return False
	return True

count6, count4 = 0, 0
for i in range(10 ** 5, 10 ** 6):
	num = [int(x) for x in str(i)]
	if alternation(str(i)) and len(num) == len(set(num)) :
		count6 += 1
for i in range(10 ** 3, 10 ** 4):
	if double(str(i)):
		count4 += 1
if count4 > count6:
	print(4, count4 - count6)
else:
	print(6, count6 - count4)

print(count6, count4)


import tensorflow as tf
from tensorflow.contrib.framework.python.ops import add_arg_scope, arg_scope
import warnings
import tensorflow_probability as tfp
import numpy as np
tfb = tfp.bijectors

def dynamic_deconv_op3d(x, W, strides=[1,2,2,2,1], padding='SAME'):

    filter_size = tf.shape(W)#list(map(int, W.get_shape()))
    xs = tf.shape(x)
    target_shape = tf.shape(x)
    #channelscale = int(int(W.get_shape()[-1])/int(x.get_shape()[-1]))
    if padding == 'SAME':
        shapescaling = tf.constant([1, strides[1], strides[2], strides[3], 1])
        print(shapescaling)
        target_shape = target_shape*shapescaling
    if padding == 'VALID':
        shapescaling = tf.constant([1, strides[1], strides[2], strides[3], 1])
        shapeoffset =  tf.constant([0, filter_size[0]-1, filter_size[1]-1, filter_size[2]-1, 1])
        #This doesnt seem to be working. Figure out at some point
        target_shape = target_shape*shapescaling
        target_shape = target_shape+shapeoffset
    return tf.nn.conv3d_transpose(x, W, target_shape, strides, padding)

def dynamic_deconv3d(name, x, shape, strides=[1,2,2,2,1], activation=tf.nn.leaky_relu):
    in_dim = x.get_shape()[-1]
    W_shape = [shape[0], shape[1], shape[2], in_dim, shape[3]]
    W = tf.get_variable(name+'_W', W_shape, tf.float32, tf.contrib.layers.xavier_initializer())
    b = tf.get_variable(name+'_b', W_shape[3], tf.float32, tf.zeros_initializer)
    conv = dynamic_deconv_op3d(x, W, strides=strides)
    return activation(tf.add(conv, b))

###Following spectral normed conv3d has been contributed by Francois Lanusse

NO_OPS = 'NO_OPS'
def _l2normalize(v, eps=1e-12):
    return v / (tf.reduce_sum(v ** 2) ** 0.5 + eps)

def spectral_normed_weight(W, u=None, num_iters=1, update_collection=None, with_sigma=False, normwt=0.7):
  # Usually num_iters = 1 will be enough
    W_shape = W.shape.as_list()
    W_reshaped = tf.reshape(W, [-1, W_shape[-1]])

    if u is None:
        u = tf.get_variable("u", [1, W_shape[-1]], initializer=tf.truncated_normal_initializer(), trainable=False)

    def power_iteration(i, u_i, v_i):
        v_ip1 = _l2normalize(tf.matmul(u_i, tf.transpose(W_reshaped)))
        u_ip1 = _l2normalize(tf.matmul(v_ip1, W_reshaped))
        return i + 1, u_ip1, v_ip1

    _, u_final, v_final = tf.while_loop(
        cond=lambda i, _1, _2: i < num_iters,
        body=power_iteration,
        loop_vars=(tf.constant(0, dtype=tf.int32),
                   u, tf.zeros(dtype=tf.float32, shape=[1, W_reshaped.shape.as_list()[0]]))
      )

    if update_collection is None:
        warnings.warn('Setting update_collection to None will make u being updated every W execution. This maybe undesirable'
                  '. Please consider using a update collection instead.')
        sigma = tf.matmul(tf.matmul(v_final, W_reshaped), tf.transpose(u_final))[0, 0]
        # sigma = tf.reduce_sum(tf.matmul(u_final, tf.transpose(W_reshaped)) * v_final)
        scaling = tf.minimum(1., normwt/sigma)
        #scaling = normwt/sigma
        W_bar = W_reshaped *scaling
        with tf.control_dependencies([u.assign(u_final)]):
            W_bar = tf.reshape(W_bar, W_shape)
    else:
        sigma = tf.matmul(tf.matmul(v_final, W_reshaped), tf.transpose(u_final))[0, 0]
        # sigma = tf.reduce_sum(tf.matmul(u_final, tf.transpose(W_reshaped)) * v_final)
        W_bar = W_reshaped / sigma
        W_bar = tf.reshape(W_bar, W_shape)
        # Put NO_OPS to not update any collection. This is useful for the second call of discriminator if the update_op
        # has already been collected on the first call.
        if update_collection != NO_OPS:
            tf.add_to_collection(update_collection, u.assign(u_final))
    if with_sigma:
        return W_bar, sigma
    else:
        return W_bar

@add_arg_scope
def scope_has_variables(scope):
    return len(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=scope.name)) > 0

@add_arg_scope
def specnormconv3d(input_, output_dim,
           kernel_size=3, stride=1, stddev=None,
           name="conv3d", spectral_normed=True, update_collection=None, with_w=False,
                   padding="SAME", reuse=tf.AUTO_REUSE, num_iters=1, normwt=1):

    k_h, k_w, k_z = [kernel_size]*3
    d_h, d_w, d_z = [stride]*3
    # Glorot intialization
    # For RELU nonlinearity, it's sqrt(2./(n_in)) instead
    fan_in = k_h * k_w * k_z * input_.get_shape().as_list()[-1]
    fan_out = k_h * k_w * k_z * output_dim
    if stddev is None:
        stddev = tf.sqrt(2. / (fan_in))

    with tf.variable_scope(name) as scope:
        if scope_has_variables(scope):
            scope.reuse_variables()
        w = tf.get_variable("w", [k_h, k_w, k_z, input_.get_shape()[-1], output_dim],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        if spectral_normed:
            conv = tf.nn.conv3d(input_, spectral_normed_weight(w, update_collection=update_collection, num_iters=num_iters, normwt=normwt),
                              strides=[1, d_h, d_w, d_z, 1], padding=padding)
        else:
            conv = tf.nn.conv3d(input_, w, strides=[1, d_h, d_w, d_z, 1], padding=padding)

        biases = tf.get_variable("b", [output_dim], initializer=tf.constant_initializer(0.0))
        conv = tf.nn.bias_add(conv, biases)
        if with_w:
            return conv, w, biases
        else:
            return conv

@add_arg_scope
def specnormconv2d(input_, output_dim,
           kernel_size=3, stride=1, stddev=None,
           name="conv2d", spectral_normed=True, update_collection=None, with_w=False,
                   padding="SAME", reuse=tf.AUTO_REUSE, num_iters=1, normwt=1.):

    k_h, k_w = [kernel_size]*2
    d_h, d_w = [stride]*2
    # Glorot intialization
    # For RELU nonlinearity, it's sqrt(2./(n_in)) instead
    fan_in = k_h * k_w * input_.get_shape().as_list()[-1]
    fan_out = k_h * k_w * output_dim
    if stddev is None:
        stddev = tf.sqrt(2. / (fan_in))

    with tf.variable_scope(name) as scope:
        if scope_has_variables(scope):
            scope.reuse_variables()
        w = tf.get_variable("w", [k_h, k_w, input_.get_shape()[-1], output_dim],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        if spectral_normed:
            conv = tf.nn.conv2d(input_, spectral_normed_weight(w, update_collection=update_collection,
                                                               num_iters=num_iters, normwt=normwt),
                              strides=[1, d_h, d_w, 1], padding=padding)
        else:
            conv = tf.nn.conv2d(input_, w, strides=[1, d_h, d_w, 1], padding=padding)

        biases = tf.get_variable("b", [output_dim], initializer=tf.constant_initializer(0.0))
        conv = tf.nn.bias_add(conv, biases)
        if with_w:
            return conv, w, biases
        else:
            return conv

@add_arg_scope
def specnormconv1d(input_, output_dim,
           kernel_size=3, stride=1, stddev=None,
           name="conv1d", spectral_normed=True, update_collection=None, with_w=False,
                   padding="SAME", reuse=tf.AUTO_REUSE, num_iters=1):

    k_h = kernel_size
    d_h = stride
    # Glorot intialization
    # For RELU nonlinearity, it's sqrt(2./(n_in)) instead
    fan_in = k_h * input_.get_shape().as_list()[-1]
    fan_out = k_h * output_dim
    if stddev is None:
        stddev = tf.sqrt(2. / (fan_in))

    with tf.variable_scope(name) as scope:
        if scope_has_variables(scope):
            scope.reuse_variables()
        w = tf.get_variable("w", [k_h, input_.get_shape()[-1], output_dim],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        #w = tf.get_variable("w", [k_h, input_.get_shape()[-1], output_dim],
        #                    initializer=tf.glorot_uniform_initializer())
        if spectral_normed:
            conv = tf.nn.conv1d(input_, spectral_normed_weight(w, update_collection=update_collection,
                                                               num_iters=num_iters),
                              stride=d_h, padding=padding)
        else:
            conv = tf.nn.conv1d(input_, w, stride=d_h, padding=padding)

        biases = tf.get_variable("b", [output_dim], initializer=tf.constant_initializer(0.0))
        conv = tf.nn.bias_add(conv, biases)
        if with_w:
            return conv, w, biases
        else:
            return conv

###Following have been taken from https://github.com/openai/glow/blob/master/tfops.py
###and modified to 3D. In addition, some variable names have been modified
###and assumptions on shape of input being constant relaxed

def int_shape(x):
    if str(x.get_shape()[0]) != '?':
        return list(map(int, x.get_shape()))
    return [-1]+list(map(int, x.get_shape()[1:]))

# Invertible 1x1 conv
@add_arg_scope
def invertible_1x1_conv(name, z, reverse=False):

    if True:  # Set to "False" to use the LU-decomposed version

        with tf.variable_scope(name, reuse=tf.AUTO_REUSE):

            shape = int_shape(z)
            w_shape = [shape[-1], shape[-1]]

            # Sample a random orthogonal matrix:
            w_init = np.linalg.qr(np.random.randn(
                *w_shape))[0].astype('float32')

            w = tf.get_variable("W", dtype=tf.float32, initializer=w_init)
            # dlogdet = tf.linalg.LinearOperator(w).log_abs_determinant() * shape[1]*shape[2]
            logdet = tf.cast(tf.log(abs(tf.matrix_determinant(
                tf.cast(w, 'float64')))), 'float32') * shape[1]*shape[2]*shape[3]
            if not reverse:

                _w = tf.reshape(w, [1, 1, 1] + w_shape)
                z = tf.nn.conv3d(z, _w, [1, 1, 1, 1, 1],
                                 'SAME', data_format='NDHWC')
                return z, logdet
            else:

                _w = tf.matrix_inverse(w)
                _w = tf.reshape(_w, [1, 1, 1]+w_shape)
                z = tf.nn.conv3d(z, _w, [1, 1, 1, 1, 1],
                                 'SAME', data_format='NDHWC')
                return z, -logdet

    else:

        # LU-decomposed version
        shape = int_shape(z)
        with tf.variable_scope(name, reuse=tf.AUTO_REUSE):

            dtype = 'float64'

            # Random orthogonal matrix:
            import scipy
            np_w = scipy.linalg.qr(np.random.randn(shape[-1], shape[-1]))[
                0].astype('float32')

            np_p, np_l, np_u = scipy.linalg.lu(np_w)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(abs(np_s))
            np_u = np.triu(np_u, k=1)

            p = tf.get_variable("P", initializer=np_p, trainable=False)
            l = tf.get_variable("L", initializer=np_l)
            sign_s = tf.get_variable(
                "sign_S", initializer=np_sign_s, trainable=False)
            log_s = tf.get_variable("log_S", initializer=np_log_s)
            # S = tf.get_variable("S", initializer=np_s)
            u = tf.get_variable("U", initializer=np_u)

            p = tf.cast(p, dtype)
            l = tf.cast(l, dtype)
            sign_s = tf.cast(sign_s, dtype)
            log_s = tf.cast(log_s, dtype)
            u = tf.cast(u, dtype)

            w_shape = [shape[-1], shape[-1]]

            l_mask = np.tril(np.ones(w_shape, dtype=dtype), -1)
            l = l * l_mask + tf.eye(*w_shape, dtype=dtype)
            u = u * np.transpose(l_mask) + tf.diag(sign_s * tf.exp(log_s))
            w = tf.matmul(p, tf.matmul(l, u))

            if True:
                u_inv = tf.matrix_inverse(u)
                l_inv = tf.matrix_inverse(l)
                p_inv = tf.matrix_inverse(p)
                w_inv = tf.matmul(u_inv, tf.matmul(l_inv, p_inv))
            else:
                w_inv = tf.matrix_inverse(w)

            w = tf.cast(w, tf.float32)
            w_inv = tf.cast(w_inv, tf.float32)
            log_s = tf.cast(log_s, tf.float32)

            if not reverse:

                w = tf.reshape(w, [1, 1, 1] + w_shape)
                z = tf.nn.conv3d(z, w, [1, 1, 1, 1, 1],
                                 'SAME', data_format='NDHWC')
                logdet = tf.reduce_sum(log_s) * (shape[1]*shape[2]*shape[3])

                return z, logdet
            else:

                w_inv = tf.reshape(w_inv, [1, 1, 1]+w_shape)
                z = tf.nn.conv3d(
                    z, w_inv, [1, 1, 1, 1, 1], 'SAME', data_format='NDHWC')
                logdet = -tf.reduce_sum(log_s) * (shape[1]*shape[2]*shape[3])

                return z, logdet

@add_arg_scope
def get_variable_ddi(name, shape, initial_value, dtype=tf.float32, init=False, trainable=True):
    w = tf.get_variable(name, shape, dtype, None, trainable=trainable)
    if init:
        w = w.assign(initial_value)
        with tf.control_dependencies([w]):
            return w
    return w

@add_arg_scope
def actnorm3d(x, scale=1., logscale_factor=3., batch_variance=False,
            reverse=False, init=False, is_training=True, scope='actnorm'):
    """
    Borrowed from https://github.com/openai/glow/blob/master/tfops.py
    """
    name = scope
    if arg_scope([get_variable_ddi], trainable=is_training):
        if not reverse:
            x = actnorm_center3d(name+"_center", x, reverse)
            x, logdet = actnorm_scale3d(name+"_scale", x, scale,
                              logscale_factor, batch_variance, reverse, init)
        else:
            x, logdet = actnorm_scale3d(name + "_scale", x, scale,
                              logscale_factor, batch_variance, reverse, init)

            x = actnorm_center3d(name+"_center", x, reverse)
        return x, logdet

@add_arg_scope
def actnorm_center3d(name, x, reverse=False):
    shape = x.get_shape()
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        assert len(shape) == 5
        x_mean = tf.reduce_mean(x, [0, 1, 2, 3], keepdims=True)
        b = get_variable_ddi(
#             "b", (1, 1, 1, 1, int_shape(x)[4]), initial_value=-x_mean)
            "b", (1, 1, 1, 1, shape[4]), initial_value=-x_mean)

        if not reverse:
            x += b
        else:
            x -= b

        return x

@add_arg_scope
def actnorm_scale3d(name, x, scale=1., logscale_factor=3.,
                  batch_variance=False, reverse=False, init=False, trainable=True):
    shape = x.get_shape()
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE), arg_scope([get_variable_ddi], trainable=trainable):
        assert len(shape) == 5
        x_var = tf.reduce_mean(x**2, [0, 1, 2, 3], keepdims=True)
#         logdet_factor = int(shape[1])*int(shape[2])*int(shape[3])
        logdet_factor = int(shape[1])*int(shape[2])*int(shape[3])
#         _shape = (1, 1, 1, 1, int_shape(x)[4])
        _shape = (1, 1, 1, 1, shape[4])

        if batch_variance:
            x_var = tf.reduce_mean(x**2, keepdims=True)

        if True:
            logs = get_variable_ddi("logs", _shape, initial_value=tf.log(
                scale/(tf.sqrt(x_var)+1e-6))/logscale_factor)*logscale_factor
            if not reverse:
                x = x * tf.exp(logs)
            else:
                x = x * tf.exp(-logs)
        else:
            # Alternative, doesn't seem to do significantly worse or better than the logarithmic version above
            s = get_variable_ddi("s", _shape, initial_value=scale /
                                 (tf.sqrt(x_var) + 1e-6) / logscale_factor)*logscale_factor
            logs = tf.log(tf.abs(s))
            if not reverse:
                x *= s
            else:
                x /= s

        dlogdet = tf.reduce_sum(logs) * logdet_factor
        if reverse:
            dlogdet *= -1
        return x, dlogdet

#Squeeze operation for invertible downsampling in 3D
#
#
class Squeeze3d(tfb.Reshape):
    """
    Borrowed from https://github.com/openai/glow/blob/master/tfops.py
    """
    def __init__(self,
                 event_shape_in,
                 factor=2,
                 is_constant_jacobian=True,
                 validate_args=False,
                 name=None):

        assert factor >= 1
        name = name or "squeeze"
        self.factor = factor
        event_shape_out = 1*event_shape_in
        event_shape_out[0] //=2
        event_shape_out[1] //=2
        event_shape_out[2] //=2
        event_shape_out[3] *=8
        self.event_shape_out = event_shape_out

        super(Squeeze3d, self).__init__(
            event_shape_out=event_shape_out,
            event_shape_in=event_shape_in,
        validate_args=validate_args,
        name=name)

    def _forward(self, x):
        if self.factor == 1:
            return x
        factor = self.factor

        shape = tf.shape(x)
        height = shape[1]
        width = shape[2]
        length = shape[3]
        n_channels = x.get_shape()[4]

#         print(height, width, length, n_channels )
#         assert height % factor == 0 and width % factor == 0 and length % factor == 0
        x = tf.reshape(x, [-1, height//factor, factor,
                           width//factor, factor, length//factor, factor, n_channels])
        x = tf.transpose(x, [0, 1, 3, 5, 7, 2, 4, 6])
        x = tf.reshape(x, [-1, height//factor, width//factor,
                               length//factor, n_channels*factor**3])
        return x

    def _inverse(self, x):
        if self.factor == 1:
            return x
        factor = self.factor

        shape = tf.shape(x)
        height = shape[1]
        width = shape[2]
        length = shape[3]
        n_channels = int(x.get_shape()[4])

#         print(height, width, length, n_channels )
        assert n_channels >= 8 and n_channels % 8 == 0
        x = tf.reshape(
            x, [-1, height, width, length, int(n_channels/factor**3), factor, factor, factor])
        x = tf.transpose(x, [0, 1, 5, 2, 6, 3, 7, 4])
        x = tf.reshape(x, (-1, height*factor,
                           width*factor, height*factor, int(n_channels/factor**3)))
        return x

def f_net(name, h, width, n_out=None):
    n_out = n_out or int(h.get_shape()[4])
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        h = tf.layers.conv3d(h, width, 3, padding='SAME', activation=tf.nn.relu)
        h = tf.layers.conv3d(h, n_out, 3, padding='SAME', activation=None)
    return h

def f_net3d(name, h, width, n_out=None):

    k_h, k_w, k_z = [3]*3
    fan_in = k_h * k_w * k_z * h.get_shape().as_list()[-1]
    stddev = tf.sqrt(2. / (fan_in))
    n_out = n_out or int(h.get_shape()[4])

    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
#         if scope_has_variables(scope):
#             scope.reuse_variables()
        w1 = tf.get_variable("w1", [3, 3, 3, h.get_shape()[-1], width],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        w2 = tf.get_variable("w2", [3, 3, 3, width, width],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        wout = tf.get_variable("wout", [3, 3, 3, width, n_out],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))

        h = tf.nn.relu(tf.nn.conv3d(h, w1, strides=[1, 1, 1, 1, 1], padding='SAME'))
        h = tf.nn.relu(tf.nn.conv3d(h, w2, strides=[1, 1, 1, 1, 1], padding='SAME'))
        h = tf.nn.conv3d(h, wout, strides=[1, 1, 1, 1, 1], padding='SAME')
    return h


import os,json
from PIL import Image as I

def loadconf(fp=os.path.join(".","config.json")):
    global conf,infp,outfp
    f=open(fp,"r+")
    conf=json.load(f)
    f.close()
    infp=conf['input']
    outfp=conf['output']

def checkdirs():
    if not os.path.isdir(infp):
        if not os.path.isfile(infp):
            raise Exception("Input Directory Does Not Exist.")
        else:
            raise Exception("Input Directory Should Not Be A File.")
    else:
        if len(os.listdir(infp))==0:
            raise Exception("Input Directory Is Empty.")
    if not os.path.isdir(outfp):
        if not os.path.isfile(outfp):
            os.mkdir(outfp)
        else:
            raise Exception("Output Directory Should Not Be A File.")

def lookdir(d,*,ext=".png"):
    if not ext.startswith("."):
        ext=f".{ext}"
    l=[]
    for f in os.listdir(d):
        if f.endswith(ext):
            l.append(f)
    return l

def convert(fn):
    img=I.open(os.path.join(infp,fn))
    o=img.crop((0,0,64,32))
    o.save(os.path.join(outfp,fn))

def loop():
    l=lookdir(infp)
    for f in l:
        convert(f)

def main():
    print("Starting...")
    loadconf()
    checkdirs()
    loop()
    print("Done!!!")

if __name__=="__main__":
    main()