# -*- coding: utf-8 -*-
"""
Y'CbCr Colour Encoding
======================

Defines the *Y'CbCr* colour encoding related transformations:

-   :func:`colour.RGB_to_YCbCr`
-   :func:`colour.YCbCr_to_RGB`
-   :func:`colour.RGB_to_YcCbcCrc`
-   :func:`colour.YcCbcCrc_to_RGB`

Notes
-----
-   *Y'CbCr* is not an absolute colourspace.

References
----------
-   :cite:`InternationalTelecommunicationUnion2011e` : International
    Telecommunication Union. (2011). Recommendation ITU-T T.871 - Information
    technology - Digital compression and coding of continuous-tone still
    images: JPEG File Interchange Format (JFIF).
    https://www.itu.int/rec/dologin_pub.asp?lang=e&\
id=T-REC-T.871-201105-I!!PDF-E&type=items
-   :cite:`InternationalTelecommunicationUnion2015h` : International
    Telecommunication Union. (2015). Recommendation ITU-R BT.2020 - Parameter
    values for ultra-high definition television systems for production and
    international programme exchange (pp. 1-8).
    https://www.itu.int/dms_pubrec/itu-r/rec/bt/\
R-REC-BT.2020-2-201510-I!!PDF-E.pdf
-   :cite:`InternationalTelecommunicationUnion2015i` : International
    Telecommunication Union. (2015). Recommendation ITU-R BT.709-6 - Parameter
    values for the HDTV standards for production and international programme
    exchange BT Series Broadcasting service (pp. 1-32).
    https://www.itu.int/dms_pubrec/itu-r/rec/bt/\
R-REC-BT.709-6-201506-I!!PDF-E.pdf
-   :cite:`SocietyofMotionPictureandTelevisionEngineers1999b` : Society of
    Motion Picture and Television Engineers. (1999). ANSI/SMPTE 240M-1995 -
    Signal Parameters - 1125-Line High-Definition Production Systems (pp. 1-7).
    http://car.france3.mars.free.fr/HD/\
INA-%2026%20jan%2006/SMPTE%20normes%20et%20confs/s240m.pdf
-   :cite:`Wikipedia2004d` : Wikipedia. (2004). YCbCr. Retrieved February 29,
    2016, from https://en.wikipedia.org/wiki/YCbCr
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.constants import DEFAULT_FLOAT_DTYPE, DEFAULT_INT_DTYPE
from colour.models.rgb.transfer_functions import (
    CV_range, eotf_inverse_BT2020, eotf_BT2020)
from colour.utilities import (CaseInsensitiveMapping, as_float_array,
                              domain_range_scale, from_range_1, to_domain_1,
                              tsplit, tstack)

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

__all__ = [
    'WEIGHTS_YCBCR', 'YCbCr_ranges', 'RGB_to_YCbCr', 'YCbCr_to_RGB',
    'RGB_to_YcCbcCrc', 'YcCbcCrc_to_RGB'
]

WEIGHTS_YCBCR = CaseInsensitiveMapping({
    'ITU-R BT.601': np.array([0.2990, 0.1140]),
    'ITU-R BT.709': np.array([0.2126, 0.0722]),
    'ITU-R BT.2020': np.array([0.2627, 0.0593]),
    'SMPTE-240M': np.array([0.2122, 0.0865])
})
"""
Luma weightings presets.

References
----------
:cite:`InternationalTelecommunicationUnion2011e`,
:cite:`InternationalTelecommunicationUnion2015i`,
:cite:`InternationalTelecommunicationUnion2015h`,
:cite:`SocietyofMotionPictureandTelevisionEngineers1999b`,
:cite:`Wikipedia2004d`

WEIGHTS_YCBCR : dict
    **{'ITU-R BT.601', 'ITU-R BT.709', 'ITU-R BT.2020', 'SMPTE-240M}**
"""

def YCbCr_ranges(bits, is_legal, is_int):
    """"
    Returns the *Y'CbCr* colour encoding ranges array for given bit depth,
    range legality and representation.

    Parameters
    ----------
    bits : int
        Bit depth of the *Y'CbCr* colour encoding ranges array.
    is_legal : bool
        Whether the *Y'CbCr* colour encoding ranges array is legal.
    is_int : bool
        Whether the *Y'CbCr* colour encoding ranges array represents integer
        code values.

    Returns
    -------
    ndarray
        *Y'CbCr* colour encoding ranges array.

    Examples
    --------
    >>> YCbCr_ranges(8, True, True)
    array([ 16, 235,  16, 240])
    >>> YCbCr_ranges(8, True, False)  # doctest: +ELLIPSIS
    array([ 0.0627451...,  0.9215686...,  0.0627451...,  0.9411764...])
    >>> YCbCr_ranges(10, False, False)
    array([ 0. ,  1. , -0.5,  0.5])
    """

    if is_legal:
        ranges = np.array([16, 235, 16, 240])
        ranges *= 2 ** (bits - 8)
    else:
        ranges = np.array([0, 2 ** bits - 1, 0, 2 ** bits - 1])

    if not is_int:
        ranges = ranges.astype(DEFAULT_FLOAT_DTYPE) / (2 ** bits - 1)

    if is_int and not is_legal:
        ranges[3] = 2 ** bits

    if not is_int and not is_legal:
        ranges[2] = -0.5
        ranges[3] = 0.5

    return ranges

def RGB_to_YCbCr(RGB,
                 K=WEIGHTS_YCBCR['ITU-R BT.709'],
                 in_bits=10,
                 in_legal=False,
                 in_int=False,
                 out_bits=8,
                 out_legal=True,
                 out_int=False,
                 **kwargs):
    """
    Converts an array of *R'G'B'* values to the corresponding *Y'CbCr* colour
    encoding values array.

    Parameters
    ----------
    RGB : array_like
        Input *R'G'B'* array of floats or integer values.
    K : array_like, optional
        Luma weighting coefficients of red and blue. See
        :attr:`colour.WEIGHTS_YCBCR` for presets. Default is
        *(0.2126, 0.0722)*, the weightings for *ITU-R BT.709*.
    in_bits : int, optional
        Bit depth for integer input, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Default is *10*.
    in_legal : bool, optional
        Whether to treat the input values as legal range. Default is *False*.
    in_int : bool, optional
        Whether to treat the input values as ``in_bits`` integer code values.
        Default is *False*.
    out_bits : int, optional
        Bit depth for integer output, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Ignored if ``out_legal`` and
        ``out_int`` are both *False*. Default is *8*.
    out_legal : bool, optional
        Whether to return legal range values. Default is *True*.
    out_int : bool, optional
        Whether to return values as ``out_bits`` integer code values. Default
        is *False*.

    Other Parameters
    ----------------
    in_range : array_like, optional
        Array overriding the computed range such as
        *in_range = (RGB_min, RGB_max)*. If ``in_range`` is undefined,
        *RGB_min* and *RGB_max* will be computed using :func:`colour.CV_range`
        definition.
    out_range : array_like, optional
        Array overriding the computed range such as
        *out_range = (Y_min, Y_max, C_min, C_max)`. If ``out_range`` is
        undefined, *Y_min*, *Y_max*, *C_min* and *C_max* will be computed
        using :func:`colour.models.rgb.ycbcr.YCbCr_ranges` definition.

    Returns
    -------
    ndarray
        *Y'CbCr* colour encoding array of integer or float values.

    Warnings
    --------
    For *Recommendation ITU-R BT.2020*, :func:`colour.RGB_to_YCbCr` definition
    is only applicable to the non-constant luminance implementation.
    :func:`colour.RGB_to_YcCbcCrc` definition should be used for the constant
    luminance case as per :cite:`InternationalTelecommunicationUnion2015h`.

    Notes
    -----

    +----------------+-----------------------+---------------+
    | **Domain \\***  | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``RGB``        | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    +----------------+-----------------------+---------------+
    | **Range \\***   | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``YCbCr``      | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    \\* This definition has input and output integer switches, thus the
    domain-range scale information is only given for the floating point mode.

    -   The default arguments, ``**{'in_bits': 10, 'in_legal': False,
        'in_int': False, 'out_bits': 8, 'out_legal': True, 'out_int': False}``
        transform a float *R'G'B'* input array normalised to domain [0, 1]
        (``in_bits`` is ignored) to a float *Y'CbCr* output array where *Y'* is
        normalised to range [16 / 255, 235 / 255] and *Cb* and *Cr* are
        normalised to range [16 / 255, 240./255]. The float values are
        calculated based on an [0, 255] integer range, but no 8-bit
        quantisation or clamping are performed.

    References
    ----------
    :cite:`InternationalTelecommunicationUnion2011e`,
    :cite:`InternationalTelecommunicationUnion2015i`,
    :cite:`SocietyofMotionPictureandTelevisionEngineers1999b`,
    :cite:`Wikipedia2004d`

    Examples
    --------
    >>> RGB = np.array([1.0, 1.0, 1.0])
    >>> RGB_to_YCbCr(RGB)  # doctest: +ELLIPSIS
    array([ 0.9215686...,  0.5019607...,  0.5019607...])

    Matching float output of The Foundry Nuke's Colorspace node set to YCbCr:

    >>> RGB_to_YCbCr(RGB,
    ...              out_range=(16 / 255, 235 / 255, 15.5 / 255, 239.5 / 255))
    ... # doctest: +ELLIPSIS
    array([ 0.9215686...,  0.5       ,  0.5       ])

    Matching float output of The Foundry Nuke's Colorspace node set to YPbPr:

    >>> RGB_to_YCbCr(RGB, out_legal=False, out_int=False)
    ... # doctest: +ELLIPSIS
    array([ 1.,  0.,  0.])

    Creating integer code values as per standard 10-bit SDI:

    >>> RGB_to_YCbCr(RGB, out_legal=True, out_bits=10, out_int=True)
    ... # doctest: +ELLIPSIS
    array([940, 512, 512]...)

    For JFIF JPEG conversion as per ITU-T T.871
    :cite:`InternationalTelecommunicationUnion2011e`:

    >>> RGB = np.array([102, 0, 51])
    >>> RGB_to_YCbCr(RGB, K=WEIGHTS_YCBCR['ITU-R BT.601'], in_range=(0, 255),
    ...              out_range=(0, 255, 0, 256), out_int=True)
    ... # doctest: +ELLIPSIS
    array([ 36, 136, 175]...)

    Note the use of 256 for the max *Cb / Cr* value, which is required so that
    the *Cb* and *Cr* output is centered about 128. Using 255 centres it
    about 127.5, meaning that there is no integer code value to represent
    achromatic colours. This does however create the possibility of output
    integer codes with value of 256, which cannot be stored in 8-bit integer
    representation. Recommendation ITU-T T.871 specifies these should be
    clamped to 255.

    These JFIF JPEG ranges are also obtained as follows:

    >>> RGB_to_YCbCr(RGB, K=WEIGHTS_YCBCR['ITU-R BT.601'], in_bits=8,
    ...              in_int=True, out_legal=False, out_int=True)
    ... # doctest: +ELLIPSIS
    array([ 36, 136, 175]...)
    """

    if in_int:
        RGB = as_float_array(RGB)
    else:
        RGB = to_domain_1(RGB)

    Kr, Kb = K
    RGB_min, RGB_max = kwargs.get('in_range',
                                  CV_range(in_bits, in_legal, in_int))
    Y_min, Y_max, C_min, C_max = kwargs.get(
        'out_range', YCbCr_ranges(out_bits, out_legal, out_int))

    RGB_float = RGB.astype(DEFAULT_FLOAT_DTYPE) - RGB_min
    RGB_float *= 1 / (RGB_max - RGB_min)
    R, G, B = tsplit(RGB_float)

    Y = Kr * R + (1 - Kr - Kb) * G + Kb * B
    Cb = 0.5 * (B - Y) / (1 - Kb)
    Cr = 0.5 * (R - Y) / (1 - Kr)
    Y *= Y_max - Y_min
    Y += Y_min
    Cb *= C_max - C_min
    Cr *= C_max - C_min
    Cb += (C_max + C_min) / 2
    Cr += (C_max + C_min) / 2

    YCbCr = tstack([Y, Cb, Cr])
    YCbCr = np.round(YCbCr).astype(
        DEFAULT_INT_DTYPE) if out_int else from_range_1(YCbCr)

    return YCbCr

def YCbCr_to_RGB(YCbCr,
                 K=WEIGHTS_YCBCR['ITU-R BT.709'],
                 in_bits=8,
                 in_legal=True,
                 in_int=False,
                 out_bits=10,
                 out_legal=False,
                 out_int=False,
                 **kwargs):
    """
    Converts an array of *Y'CbCr* colour encoding values to the corresponding
    *R'G'B'* values array.

    Parameters
    ----------
    YCbCr : array_like
        Input *Y'CbCr* colour encoding array of integer or float values.
    K : array_like, optional
        Luma weighting coefficients of red and blue. See
        :attr:`colour.WEIGHTS_YCBCR` for presets. Default is
        *(0.2126, 0.0722)*, the weightings for *ITU-R BT.709*.
    in_bits : int, optional
        Bit depth for integer input, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Default is *8*.
    in_legal : bool, optional
        Whether to treat the input values as legal range. Default is *True*.
    in_int : bool, optional
        Whether to treat the input values as ``in_bits`` integer code values.
        Default is *False*.
    out_bits : int, optional
        Bit depth for integer output, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Ignored if ``out_legal`` and
        ``out_int`` are both *False*. Default is *10*.
    out_legal : bool, optional
        Whether to return legal range values. Default is *False*.
    out_int : bool, optional
        Whether to return values as ``out_bits`` integer code values. Default
        is *False*.

    Other Parameters
    ----------------
    in_range : array_like, optional
        Array overriding the computed range such as
        *in_range = (Y_min, Y_max, C_min, C_max)*. If ``in_range`` is
        undefined, *Y_min*, *Y_max*, *C_min* and *C_max* will be computed using
        :func:`colour.models.rgb.ycbcr.YCbCr_ranges` definition.
    out_range : array_like, optional
        Array overriding the computed range such as
        *out_range = (RGB_min, RGB_max)*. If ``out_range`` is undefined,
        *RGB_min* and *RGB_max* will be computed using :func:`colour.CV_range`
        definition.

    Returns
    -------
    ndarray
        *R'G'B'* array of integer or float values.

    Notes
    -----

    +----------------+-----------------------+---------------+
    | **Domain \\***  | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``YCbCr``      | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    +----------------+-----------------------+---------------+
    | **Range \\***   | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``RGB``        | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    \\* This definition has input and output integer switches, thus the
    domain-range scale information is only given for the floating point mode.

    Warnings
    --------
    For *Recommendation ITU-R BT.2020*, :func:`colour.YCbCr_to_RGB`
    definition is only applicable to the non-constant luminance implementation.
    :func:`colour.YcCbcCrc_to_RGB` definition should be used for the constant
    luminance case as per :cite:`InternationalTelecommunicationUnion2015h`.

    References
    ----------
    :cite:`InternationalTelecommunicationUnion2011e`,
    :cite:`InternationalTelecommunicationUnion2015i`,
    :cite:`SocietyofMotionPictureandTelevisionEngineers1999b`,
    :cite:`Wikipedia2004d`

    Examples
    --------
    >>> YCbCr = np.array([502, 512, 512])
    >>> YCbCr_to_RGB(YCbCr, in_bits=10, in_legal=True, in_int=True)
    array([ 0.5,  0.5,  0.5])
    """

    if in_int:
        YCbCr = as_float_array(YCbCr)
    else:
        YCbCr = to_domain_1(YCbCr)

    Y, Cb, Cr = tsplit(YCbCr.astype(DEFAULT_FLOAT_DTYPE))
    Kr, Kb = K
    Y_min, Y_max, C_min, C_max = kwargs.get(
        'in_range', YCbCr_ranges(in_bits, in_legal, in_int))
    RGB_min, RGB_max = kwargs.get('out_range',
                                  CV_range(out_bits, out_legal, out_int))

    Y -= Y_min
    Cb -= (C_max + C_min) / 2
    Cr -= (C_max + C_min) / 2
    Y *= 1 / (Y_max - Y_min)
    Cb *= 1 / (C_max - C_min)
    Cr *= 1 / (C_max - C_min)
    R = Y + (2 - 2 * Kr) * Cr
    B = Y + (2 - 2 * Kb) * Cb
    G = (Y - Kr * R - Kb * B) / (1 - Kr - Kb)

    RGB = tstack([R, G, B])
    RGB *= RGB_max - RGB_min
    RGB += RGB_min
    RGB = np.round(RGB).astype(DEFAULT_INT_DTYPE) if out_int else from_range_1(
        RGB)

    return RGB

def RGB_to_YcCbcCrc(RGB,
                    out_bits=10,
                    out_legal=True,
                    out_int=False,
                    is_12_bits_system=False,
                    **kwargs):
    """
    Converts an array of *RGB* linear values to the corresponding *Yc'Cbc'Crc'*
    colour encoding values array.

    Parameters
    ----------
    RGB : array_like
        Input *RGB* array of linear float values.
    out_bits : int, optional
        Bit depth for integer output, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Ignored if ``out_legal`` and
        ``out_int`` are both *False*. Default is *10*.
    out_legal : bool, optional
        Whether to return legal range values. Default is *True*.
    out_int : bool, optional
        Whether to return values as ``out_bits`` integer code values. Default
        is *False*.
    is_12_bits_system : bool, optional
        *Recommendation ITU-R BT.2020* OETF (OECF) adopts different parameters
        for 10 and 12 bit systems. Default is *False*.

    Other Parameters
    ----------------
    out_range : array_like, optional
        Array overriding the computed range such as
        *out_range = (Y_min, Y_max, C_min, C_max)*. If ``out_range`` is
        undefined, *Y_min*, *Y_max*, *C_min* and *C_max* will be computed
        using :func:`colour.models.rgb.ycbcr.YCbCr_ranges` definition.

    Returns
    -------
    ndarray
        *Yc'Cbc'Crc'* colour encoding array of integer or float values.

    Notes
    -----

    +----------------+-----------------------+---------------+
    | **Domain \\***  | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``RGB``        | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    +----------------+-----------------------+---------------+
    | **Range \\***   | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``YcCbcCrc``   | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    \\* This definition has input and output integer switches, thus the
    domain-range scale information is only given for the floating point mode.

    Warnings
    --------
    This definition is specifically for usage with
    *Recommendation ITU-R BT.2020* when adopting the constant luminance
    implementation.

    References
    ----------
    :cite:`InternationalTelecommunicationUnion2015h`, :cite:`Wikipedia2004d`

    Examples
    --------
    >>> RGB = np.array([0.18, 0.18, 0.18])
    >>> RGB_to_YcCbcCrc(RGB, out_legal=True, out_bits=10, out_int=True,
    ...                 is_12_bits_system=False)
    ... # doctest: +ELLIPSIS
    array([422, 512, 512]...)
    """

    R, G, B = tsplit(to_domain_1(RGB))
    Y_min, Y_max, C_min, C_max = kwargs.get(
        'out_range', YCbCr_ranges(out_bits, out_legal, out_int))

    Yc = 0.2627 * R + 0.6780 * G + 0.0593 * B

    with domain_range_scale('ignore'):
        Yc = eotf_inverse_BT2020(Yc, is_12_bits_system=is_12_bits_system)
        R = eotf_inverse_BT2020(R, is_12_bits_system=is_12_bits_system)
        B = eotf_inverse_BT2020(B, is_12_bits_system=is_12_bits_system)

    Cbc = np.where((B - Yc) <= 0, (B - Yc) / 1.9404, (B - Yc) / 1.5816)
    Crc = np.where((R - Yc) <= 0, (R - Yc) / 1.7184, (R - Yc) / 0.9936)
    Yc *= Y_max - Y_min
    Yc += Y_min
    Cbc *= C_max - C_min
    Crc *= C_max - C_min
    Cbc += (C_max + C_min) / 2
    Crc += (C_max + C_min) / 2

    YcCbcCrc = tstack([Yc, Cbc, Crc])
    YcCbcCrc = (np.round(YcCbcCrc).astype(DEFAULT_INT_DTYPE)
                if out_int else from_range_1(YcCbcCrc))

    return YcCbcCrc

def YcCbcCrc_to_RGB(YcCbcCrc,
                    in_bits=10,
                    in_legal=True,
                    in_int=False,
                    is_12_bits_system=False,
                    **kwargs):
    """
    Converts an array of *Yc'Cbc'Crc'* colour encoding values to the
    corresponding *RGB* array of linear values.

    Parameters
    ----------
    YcCbcCrc : array_like
        Input *Yc'Cbc'Crc'* colour encoding array of linear float values.
    in_bits : int, optional
        Bit depth for integer input, or used in the calculation of the
        denominator for legal range float values, i.e. 8-bit means the float
        value for legal white is *235 / 255*. Default is *10*.
    in_legal : bool, optional
        Whether to treat the input values as legal range. Default is *False*.
    in_int : bool, optional
        Whether to treat the input values as ``in_bits`` integer code values.
        Default is *False*.
    is_12_bits_system : bool, optional
        *Recommendation ITU-R BT.2020* EOTF (EOCF) adopts different parameters
        for 10 and 12 bit systems. Default is *False*.

    Other Parameters
    ----------------
    in_range : array_like, optional
        Array overriding the computed range such as
        *in_range = (Y_min, Y_max, C_min, C_max)*. If ``in_range`` is
        undefined, *Y_min*, *Y_max*, *C_min* and *C_max* will be computed using
        :func:`colour.models.rgb.ycbcr.YCbCr_ranges` definition.

    Returns
    -------
    ndarray
        *RGB* array of linear float values.

    Notes
    -----

    +----------------+-----------------------+---------------+
    | **Domain \\***  | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``YcCbcCrc``   | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    +----------------+-----------------------+---------------+
    | **Range \\***   | **Scale - Reference** | **Scale - 1** |
    +================+=======================+===============+
    | ``RGB``        | [0, 1]                | [0, 1]        |
    +----------------+-----------------------+---------------+

    \\* This definition has input and output integer switches, thus the
    domain-range scale information is only given for the floating point mode.

    Warnings
    --------
    This definition is specifically for usage with
    *Recommendation ITU-R BT.2020* when adopting the constant luminance
    implementation.

    References
    ----------
    :cite:`InternationalTelecommunicationUnion2015h`,
    :cite:`Wikipedia2004d`

    Examples
    --------
    >>> YcCbcCrc = np.array([1689, 2048, 2048])
    >>> YcCbcCrc_to_RGB(YcCbcCrc, in_legal=True, in_bits=12, in_int=True,
    ...                 is_12_bits_system=True)
    ... # doctest: +ELLIPSIS
    array([ 0.1800903...,  0.1800903...,  0.1800903...])
    """

    if in_int:
        YcCbcCrc = as_float_array(YcCbcCrc)
    else:
        YcCbcCrc = to_domain_1(YcCbcCrc)

    Yc, Cbc, Crc = tsplit(YcCbcCrc.astype(DEFAULT_FLOAT_DTYPE))
    Y_min, Y_max, C_min, C_max = kwargs.get(
        'in_range', YCbCr_ranges(in_bits, in_legal, in_int))

    Yc -= Y_min
    Cbc -= (C_max + C_min) / 2
    Crc -= (C_max + C_min) / 2
    Yc *= 1 / (Y_max - Y_min)
    Cbc *= 1 / (C_max - C_min)
    Crc *= 1 / (C_max - C_min)
    B = np.where(Cbc <= 0, Cbc * 1.9404 + Yc, Cbc * 1.5816 + Yc)
    R = np.where(Crc <= 0, Crc * 1.7184 + Yc, Crc * 0.9936 + Yc)

    with domain_range_scale('ignore'):
        Yc = eotf_BT2020(Yc, is_12_bits_system=is_12_bits_system)
        B = eotf_BT2020(B, is_12_bits_system=is_12_bits_system)
        R = eotf_BT2020(R, is_12_bits_system=is_12_bits_system)

    G = (Yc - 0.0593 * B - 0.2627 * R) / 0.6780

    RGB = tstack([R, G, B])

    return from_range_1(RGB)


#!/usr/bin/python

"""
spatial weights module for Space-Time Analysis of Regional Systems
----------------------------------------------------------------------
AUTHOR(S):  Sergio J. Rey sjrey@users.sourceforge.net
----------------------------------------------------------------------

OVERVIEW:

This module implements spatial weights classes for STARS.

XXX PySAL

"""

import time
from shapereader import *
import numpy.oldnumeric as Numeric
import numpy.oldnumeric.linear_algebra as LinearAlgebra

import Gis # for centroids on disconnected island polygons

# delta to get buckets right
DELTA = 0.000001
# constants for weights types
WT_UNKNOWN = 0
WT_ROOK = 1
WT_QUEEN = 2
# constants for bucket sizes
BUCK_SM = 8
BUCK_LG = 20
SHP_SMALL = 1000
# constant for spatial weights size
WT_SMALL = 1000
class spweight:
    """
    spweight data structure: list of lists
    [ metadata ], [key:obs dictionary] [obs:key dictionary]
    [number of neighbors] [neighbor id lists ] [neighbor weights lists]
    [characteristics ] [traces ] [eigvalues]
    need a dbf reader to use a "key" variable instead of sequence number
    sequence number is fine as long as "pure" shape files since the matching order
    is part of the ESRI shape file format
    when called without filename, initialize only, useful for grid cases
    """

    def __init__(self,filename="",wtType=WT_ROOK):
    # initialization here: check input file type and
    # call appropriate reader
        self.meta = []
        self.keyobs = {}
        self.obskey = {}
        self.numneigh = []       # number of neighbors
        self.neighbors = []      # neighbor ids
        self.weights  = []
        self.characteristics = []
        self.traces = []
        self.eigvalues = []
        if filename:
            # only shape input implemented so far
            self.shp2wt(filename,wtType)
        else:   # must create shapefile from scratch grid only
            pass
        self.numneighbrs()    # computes the number of neighbors

################### constructors

    # read gal file and convert to wt data structure
    def wtfromgal(self):
        pass

    # read gwt file and convert to wt data structure
    def wtfromgwt(self):
        pass
#------------------
    # read shp file and construct rook or queen contiguity
    def shp2wt(self,filename,wtType):
        raw_shape = shapefile(filename)
        if raw_shape.shptype == SHP_POINT:
            return       # need handler
        shapepoints = raw_shape.shplist    # list of lists
        shapebox = raw_shape.shpbox      # bounding box
        self.shapepoints = shapepoints

        numPoly = len(shapepoints)
        # bucket size
        if (numPoly < SHP_SMALL):
            bucketMin = numPoly / BUCK_SM + 2
        else:
            bucketMin = numPoly / BUCK_LG + 2
        # bucket length
        lengthX = ((shapebox[2]+DELTA) - shapebox[0]) / bucketMin
        lengthY = ((shapebox[3]+DELTA) - shapebox[1]) / bucketMin

        # initialize buckets
        bucketX = [ [] for i in range(bucketMin) ]
        bucketY = [ [] for i in range(bucketMin) ]
        polyXbucket = [ [] for i in range(numPoly) ]  # list with buckets for X
        polyYbucket = [ [] for i in range(numPoly) ]  # list with buckets for Y
        self.neighbors = [ [] for i in range(numPoly) ]      # list of lists for neighbors

        minbox = shapebox[:2] * 2  # minX,minY,minX,minY
        blen = [lengthX,lengthY] * 2  # lenx,leny,lenx,leny

        for i in range(numPoly):
            pb = [int((shapepoints[i][0][j] - minbox[j])/blen[j]) for j in range(4)]
            for j in range(pb[0],pb[2]+1):
                polyXbucket[i].append(j)
                bucketX[j].append(i)
            for j in range(pb[1],pb[3]+1):
                polyYbucket[i].append(j)
                bucketY[j].append(i)

        #create candidate neighbors from buckets
        for i in range(numPoly):
            buckX = []
            for j in range(0,len(polyXbucket[i])):
                buckX += bucketX[polyXbucket[i][j]]
            buckY = []
            for j in range(0,len(polyYbucket[i])):
                buckY += bucketY[polyYbucket[i][j]]
            buckX = dict( [ (j,j) for j in buckX ]).keys()
            buckY = dict( [ (j,j) for j in buckY ]).keys()
            buckX.sort()
            buckY.sort()

            nb = []
            if len(buckX) < len(buckY):
                k = buckX.index(i) + 1
                nb = [ buckX[jj] for jj in range(k,len(buckX))
                        if (buckX[jj] in buckY) and (shapepoints[i][0].bbcommon(shapepoints[buckX[jj]][0]) ) ]
            else:
                k = buckY.index(i) + 1
                nb = [ buckY[jj] for jj in range(k,len(buckY)) if  (buckY[jj] in buckX)
                        and (shapepoints[i][0].bbcommon(shapepoints[buckY[jj]][0])) ]
            for ii in range(0,len(nb)):
                ch=0
                jj=0
                kk=len(shapepoints[i][2]) - 1
                nbi = nb[ii]

                while not ch and jj < kk:
                    if wtType == WT_ROOK:
                        ch = (shapepoints[i][2][jj] in shapepoints[nbi][2]) and (shapepoints[i][2][jj+1] in shapepoints[nbi][2])
                    else:   # queen
                        ch = shapepoints[i][2][jj] in shapepoints[nbi][2]
                    jj += 1
                if ch:
                    self.neighbors[i].append(nbi)
                    self.neighbors[nbi].append(i)
                #print 'origin', i+1
                #print 'neighbors: ', [ nid+1 for nid in self.neighbors[i] ]

#----------------
    # construct rook for grid
    def grid2rk(self,nrows,ncols):
        m = nrows * ncols
        self.neighbors= [ [] for i in range(m) ]  # initialize
        for i in range(nrows):
            for j in range(ncols):
                k = i * ncols + j
                # to left
                if j:
                    self.neighbors[k].append(k-1)
                # to right
                if j < ncols - 1:
                    self.neighbors[k].append(k+1)
                # above
                if i:
                    self.neighbors[k].append(k - ncols)
                # below
                if i < nrows - 1:
                    self.neighbors[k].append(k + ncols)
                self.neighbors[k].sort()
        self.numneighbrs()
#-----------------
    # construct rook - torus for grid
    def grid2rktor(self,nrows,ncols):
        m = nrows * ncols
        self.neighbors= [ [] for i in range(m) ]  # initialize
        for i in range(nrows):
            for j in range(ncols):
                k = i * ncols + j
                # to left
                if j:
                    self.neighbors[k].append(k-1)
                else:
                    self.neighbors[k].append(k + ncols -1)
                # to right
                if j < ncols - 1:
                    self.neighbors[k].append(k+1)
                else:
                    self.neighbors[k].append(k - ncols + 1)
                # above
                if i:
                    self.neighbors[k].append(k - ncols)
                else:
                    self.neighbors[k].append((nrows-1)*ncols + j)
                # below
                if i < nrows - 1:
                    self.neighbors[k].append(k + ncols)
                else:
                    self.neighbors[k].append(j)
                self.neighbors[k].sort()
        self.numneighbrs()

######################### read/write methods
    # read pickled weight file ?
    def readwt(self):
        pass

    # write wt to gal file
    def wt2gal(self):
        pass

    # write wt to gwt file
    def wt2gwt(self):
        pass

    # pickle wt file
    def wt2pickle(self):
        pass

    # conversion from wt data structure to numpy weight matrix
    # gal form only so far
    def wt2mat(self):
        n = len(self.neighbors)
        w = Numeric.zeros((n,n),Numeric.Float)
        for i in range(n):
            if self.numneigh[i]:
                kk = 1.0 / self.numneigh[i]
                for j in self.neighbors[i]:
                    w[i][j] = kk
        return w

    # write numpy to gal file
    def mat2gal(self):
        pass

    # write numpty to gwt file
    def mat2gwt(self):
        pass

###################### weights computations
    # number of neighbors
    def numneighbrs(self):
        self.numneigh = []   # always reinitialize
        for i in range(len(self.neighbors)):
            self.numneigh.append(len(self.neighbors[i]))

    # weights characteristics
    def wtchars(self):
        pass

    # weights traces
    def wttraces(self):
        pass

    # higher order weights
    def wt2higher(self):
        pass

    # weights eigenvalues
    def wteigen(self):
        pass

    # neighbor histogram
    def histogram(self):
        maxn = max(self.numneigh)
        counts = [self.numneigh.count(i) for i in range(maxn+1) ]
        return counts

    def islands(self):
        """returns ids of any island observations."""
        ij = zip(range(len(self.numneigh)),self.numneigh)
        return [ i for i,j in ij if j==0 ]

    def findNearestNeighbors(self, ids):
        """find nearest neighbors based on bounding box centroid distances.

        ids: list of shape ids to find neighbors for.
        """
        #bb centroids
        cent = [ ((s[0][0]+s[0][2])/2.,(s[0][1]+s[0][3])/2.) for s in
                self.shapepoints]
        self.centroids = cent
        c=cent
        neighbors = []
        rn=range(len(self.shapepoints))
        for id in ids:
            x0,y0=cent[id]
            di = [ (x0-c[j][0])*(x0-c[j][0])+(y0-c[j][1])*(y0-c[j][1])
                  for j in rn ]
            maxd=max(di)
            di[id]=maxd
            nid = di.index(min(di))
            neighbors.append(nid)
        return neighbors

    def findIslandNeighbors(self):
        islands = self.islands()
        return zip(islands,self.findNearestNeighbors(islands))

    def fixIslands(self):
        """attaches island shapes to nearest neighbor and reconstructs gal
        information accordingly"""

        islandInfo = self.findIslandNeighbors()
        addedJoins = []
        if islandInfo:
            for i,j in islandInfo:
                ineigh = self.neighbors[i]
                jneigh = self.neighbors[j]
                if j not in ineigh:
                    ineigh.append(j)
                    ineigh.sort()
                    self.numneigh[i] = len(ineigh)
                    self.neighbors[i] = ineigh
                    addedJoins.append((i,j))
                if i not in jneigh:
                    jneigh.append(i)
                    jneigh.sort()
                    self.numneigh[j] = len(jneigh)
                    self.neighbors[j] = jneigh
                    addedJoins.append((j,i))
        self.islandInfo = islandInfo
        self.addedJoins = addedJoins

    # distance weights

    # spatial lag for single variable
    # checks what is passed and initializes same
    # assumes gal form only so far
    def splag(self,x):
        n = len(x)
        if type(x) == list:
            wx = [ 0 for i in range(n) ]
        elif type(x) == type(Numeric.array(1)):
            wx = Numeric.zeros(n,Numeric.Float)
        for i in range(n):
            if self.numneigh[i]:
                for j in self.neighbors[i]:
                    wx[i] += x[j]
                wx[i] /= self.numneigh[i]
        return wx

    # spatial filter

    # spatial AR transformation
    # force = 0 default let size determine
    # force = 1 always iterative
    # precis = DELTA precision criterion same as for buckets
    # precis set precision criterion
    def sartran(self,rho,x,force=0,precis=DELTA):
        n = len(x)
        listflag = 0
        if type(x) == list:
            x = Numeric.array(x,Numeric.Float)
            listflag = 1
        sarx = Numeric.zeros(n,Numeric.Float)
        if n > WT_SMALL or force:
            sarx = x
            wx = self.splag(x) * rho
            sarx += wx
            while max(wx) > precis:
                wx = self.splag(wx) * rho
                sarx += wx
        else:   # small weights full matrix inverse
            w = self.wt2mat()
            w *= - rho
            w += Numeric.identity(n)
            wx = LinearAlgebra.inverse(w)
            sarx = Numeric.matrixmultiply(wx,x)
        if listflag:
            return sarx.tolist()
        else:
            return sarx

    # editing weights

    # visualizing the  structure of weights

    # weights computations: addition, subtraction, multiplication

# alternatively use subclasses with special constructors

###################
###################
if __name__ == "__main__":
    fname = raw_input("Enter the shape file name (include .shp): ")
    t0 = time.time()
    w1=spweight(fname,1)
    t1 = time.time()
    print "-------------------------------------"
    print "using "+str(fname)
    print "time elapsed for rook: " + str(t1-t0)
    #t2 = time.time()
    #w2 = spweight(fname,2)
    #t3 = time.time()
    #print "time elapsed for queen: " + str(t3-t2)
    #print "-------- spatial transformation ------------"
    #n = raw_input("Enter the dimension of the grid: ")
    #n = int(n)
    #n2 = n*n
    #rho = raw_input("Enter the value for rho: ")
    #rho = float(rho)
    #ff = raw_input("Force iterative procedure 1 yes 0 no: ")
    #ff = int(ff)
    #x = range(n2)
    #xx = Numeric.array(x,Numeric.Float)
    #w3 = spweight()
    #w3.grid2rk(n,n)
    #t4 = time.time()
    #wx = w3.sartran(rho,xx,force=ff,precis=0.0000001)
    #t5 = time.time()
    #print "time elapsed for sartran " + str(n) + "  by " + str(n) + " : " + str(t5-t4)


import os

import numpy as np
from scipy.constants import c as SPEED_OF_LIGHT
from simphony.elements import Model
from simphony.tools import freq2wl, interpolate, wl2freq

class ebeam_bdc_te1550(Model):
    """
    A bidirectional coupler optimized for TE polarized light at 1550 nanometers.

    The bidirectional coupler has 4 ports, labeled as pictured. Its efficiently
    splits light that is input from one port into the two outputs on the opposite
    side (with a corresponding pi/2 phase shift). Additionally, it efficiently
    interferes lights from two adjacent inputs, efficiently splitting the
    interfered signal between the two ports on the opposing side.

    .. image:: /user/libraries/images/ebeam_bdc_te1550.png
        :alt: ebeam_bdc_te1550.png
    """

    pins = ("n1", "n2", "n3", "n4")  #: The default pin names of the device
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "sparams",
            "ebeam_bdc_te1550.npz",
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

class ebeam_dc_halfring_te1550(Model):
    pins = (
        "n1",
        "n2",
    )  #: The default pin names of the device
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "sparams",
            "ebeam_dc_halfring_te1550.npz",
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

class ebeam_gc_te1550(Model):
    """
    A grating coupler optimized for TE polarized light at 1550 nanometers.

    The grating coupler efficiently couples light from a fiber array positioned
    above the chip into the circuit. For the TE mode, the angle is -25 degrees
    [needs citation].

    .. image:: /user/libraries/images/ebeam_gc_te1550.png
        :alt: ebeam_bdc_te1550.png
    """

    pins = (
        "n1",
        "n2",
    )  #: The default pin names of the device
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "sparams",
            "ebeam_gc_te1550.npz",
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

class ebeam_terminator_te1550(Model):
    """
    A terminator component that dissipates light into free space optimized for
    TE polarized light at 1550 nanometers.

    The terminator dissipates excess light into free space. If you have a path
    where the light doesn't need to be measured but you don't want it reflecting
    back into the circuit, you can use a terminator to release it from the circuit.

    .. image:: /user/libraries/images/ebeam_terminator_te1550.png
        :alt: ebeam_bdc_te1550.png
    """

    pins = ("n1",)  #: The default pin names of the device
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "sparams",
            "ebeam_terminator_te1550.npz",
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

class ebeam_wg_integral_1550(Model):
    """
    Model for an waveguide optimized for TE polarized light at 1550 nanometers.

    A waveguide easily connects other optical components within a circuit.

    .. image:: /user/libraries/images/ebeam_wg_integral_1550.png
        :alt: ebeam_bdc_te1550.png

    Parameters
    ----------
    length : float
        Waveguide length in meters.
    lam0 : float, optional
        Central wavelength for calculation in meters (default 1.55 microns).
    ne : float, optional
        Effective index (default 2.44553).
    ng : float, optional
        Group velocity (default 4.19088).
    nd : float, optional
        Group dispersion (default 3.54275e-04).
    sigma_ne : float, optional
        Standard deviation of the effective index (default 0.05).
    sigma_ng : float, optional
        Standard deviation of the group velocity (default 0.05).
    sigma_nd : float, optional
        Standard deviation of the group dispersion (default 0.0001).

    Notes
    -----
    The ``sigma_`` values in the parameters are used for monte carlo simulations.
    """

    pins = (
        "n1",
        "n2",
    )  #: The default pin names of the device
    freq_range = (
        187370000000000.0,
        199862000000000.0,
    )  #: The valid frequency range for this model.

    def __init__(
        self,
        length,
        lam0=1.55e-06,
        ne=2.44553,
        ng=4.19088,
        nd=0.000354275,
        sigma_ne=0.05,
        sigma_ng=0.05,
        sigma_nd=0.0001,
    ):
        self.length = length
        self.lam0 = lam0
        self.ne = ne
        self.ng = ng
        self.nd = nd
        self.sigma_ne = sigma_ne
        self.sigma_ng = sigma_ng
        self.sigma_nd = sigma_nd
        self.regenerate_monte_carlo_parameters()

    def s_parameters(self, freq):
        """Get the s-parameters of a waveguide.

        Parameters
        ----------
        start : float
            The starting frequency to obtain s-parameters for (in Hz).
        stop : float
            The ending frequency to obtain s-parameters for (in Hz).
        num : int
            The number of points to use between start_freq and stop_freq.

        Returns
        -------
        (frequency, s) : tuple
            Returns a tuple containing the frequency array, ``frequency``,
            corresponding to the calculated s-parameter matrix, ``s``.
        """
        return self.cacl_s_params(
            freq, self.length, self.lam0, self.ne, self.ng, self.nd
        )

    def monte_carlo_s_parameters(self, freq):
        """
        Returns a monte carlo (randomized) set of s-parameters.

        In this implementation of the monte carlo routine, random values are
        generated for ne, ng, and nd for each run through of the monte carlo
        simulation. This means that all waveguide elements throughout a single
        circuit will have the same (random) ne, ng, and nd values. Hence, there
        is correlated randomness in the monte carlo parameters but they are
        consistent within a single circuit.
        """
        return self.cacl_s_params(
            freq, self.length, self.lam0, self.rand_ne, self.rand_ng, self.rand_nd
        )

    def regenerate_monte_carlo_parameters(self):
        self.rand_ne = np.random.normal(self.ne, self.sigma_ne)
        self.rand_ng = np.random.normal(self.ng, self.sigma_ng)
        self.rand_nd = np.random.normal(self.nd, self.sigma_nd)

    @staticmethod
    def cacl_s_params(frequency, length, lam0, ne, ng, nd):
        # Initialize array to hold s-params
        s = np.zeros((len(frequency), 2, 2), dtype=complex)

        # Loss calculation
        TE_loss = 700  # dB/m for width 500nm
        alpha = TE_loss / (20 * np.log10(np.exp(1)))

        w = np.asarray(frequency) * 2 * np.pi  # get angular frequency from frequency
        w0 = (2 * np.pi * SPEED_OF_LIGHT) / lam0  # center frequency (angular)

        # calculation of K
        K = (
            2 * np.pi * ne / lam0
            + (ng / SPEED_OF_LIGHT) * (w - w0)
            - (nd * lam0 ** 2 / (4 * np.pi * SPEED_OF_LIGHT)) * ((w - w0) ** 2)
        )

        for x in range(0, len(frequency)):  # build s-matrix from K and waveguide length
            s[x, 0, 1] = s[x, 1, 0] = np.exp(-alpha * length + (K[x] * length * 1j))

        return s

class ebeam_y_1550(Model):
    """
    The y-branch efficiently splits the input between the two outputs.

    .. image:: /user/libraries/images/ebeam_y_1550.png
        :alt: ebeam_bdc_te1550.png
    """

    pins = ("n1", "n2", "n3")  #: The default pin names of the device
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)), "sparams", "ebeam_y_1550.npz"
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

class ebeam_dc_te1550(Model):
    """
    A directional coupler optimized for TE polarized light at 1550 nanometers.

    The directional coupler has 4 ports, labeled as pictured. Its efficiently
    splits light that is input from one port into the two outputs on the opposite
    side (with a corresponding pi/2 phase shift). Additionally, it efficiently
    interferes lights from two adjacent inputs, efficiently splitting the
    interfered signal between the two ports on the opposing side.

    .. image:: /user/libraries/images/ebeam_bdc_te1550.png
        :alt: ebeam_bdc_te1550.png
    """

    pins = ("n1", "n2", "n3", "n4")
    loaded = np.load(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "sparams",
            "ebeam_dc_te1550.npz",
        )
    )
    s_params = (loaded["f"], loaded["s"])
    freq_range = (
        s_params[0][0],
        s_params[0][-1],
    )  #: The valid frequency range for this model.

    def s_parameters(self, freq):
        return interpolate(freq, self.s_params[0], self.s_params[1])

if __name__ == "__main__":
    import matplotlib.pyplot as plt

    bdc = ebeam_bdc_te1550()
    wav = np.linspace(1520, 1570, 1024) * 1e-9
    f = 3e8 / wav
    s = bdc.s_parameters(freq=f)
    plt.plot(wav, np.abs(s[:, 1] ** 2))

    plt.show()


"""**This module simulates Urban Energy Requirements**.

Outputs are stored stored as csv.
"""
import pandas as pd
import geopandas as gpd
import numpy as np
import matplotlib.pyplot as plt
import os
from pathlib import Path
import pickle
import logging

logging.basicConfig(format='%(asctime)s: %(levelname)s: %(message)s',
                    filename="../code/log/demand_supply.log",
                    level=logging.DEBUG)

class UrbanEnergyRequirement:
    """Simulate urban energy requirements.

    :returns: csv files of simulated load profile and PV requirements
    """

    def __init__(self):
        """Init function."""
        self.input_destination_1 = "../data/01_raw_input_data/"
        self.input_destination_2 = "../data/02_urban_output_data/"
        self.output_destination = "../data/03_urban_energy_requirements/"
        self.output_destination2 = "../data/04_Visualisation/"
        self.norm = 1000  # standard unit converter
        if Path(self.output_destination).exists():
            pass
        else:
            os.mkdir(self.output_destination)

    def feedin_data(self):
        """Get wind and pv feed in data."""
        with open(os.path.join(self.input_destination_1, 'fp'), 'rb') as fp:
            feedin_parameter = pickle.load(fp)
        self.peak_power = feedin_parameter[0]  # in watts
        self.module_area = feedin_parameter[1]
        self.norminal_power_wind = feedin_parameter[2]  # in watts
        pv_feedin = pd.read_csv(os.path.join(
            self.input_destination_1, 'pv_power.csv'))
        wind_feedin = pd.read_csv(os.path.join(
            self.input_destination_1, 'wind_power.csv'))
        pv_feedin = pv_feedin.reindex()
        self.time_stamp = pv_feedin[['time']]
        self.pv_data = pv_feedin["pv"]
        self.wind_power_data = wind_feedin["wind"]
        logging.info("Feedin data and powersystem parameter imported.")
        print('*** Feedin data imported ***')

    def standard_load_profiles(self):
        """Read Standard Load Profiles."""
        self.dfs = pd.read_csv(os.path.join(
            self.input_destination_1, 'SLP.csv'))
        # get different scenario load profiles in kWh
        self.AL, self.CL = self.dfs['L0']/self.norm, self.dfs['G0']/self.norm
        self.EL, self.IL = self.dfs['G1']/self.norm, self.dfs['G3']/self.norm
        self.RL, self.SL = self.dfs['H0']/self.norm, self.dfs['SB2']/self.norm
        logging.info("Get load profile for different scenarios.")
        print('Info: Load profile data imported')

    def electricity_usage_index(self):
        """Calculate Electricity Usage Index for different building type kWh/m2 per year."""
        self.x_a, self.x_c = 120/self.norm, 201/self.norm
        self.x_e, self.x_i = 142/self.norm, 645/self.norm
        self.x_r, self.x_sl = 146/self.norm, 4/self.norm
        logging.info(
            "Calculate electricty usage index for urban infrastructure.")
        print('Info: Electricity usage index calculated')

    def roofTop_pv_capacity(self):
        """Calculate maximum installed pv capacity as a function of rooftop area."""
        dfa = gpd.read_file(os.path.join(
            self.input_destination_2, 'agricultural/agricultural.shp'))
        dfc = gpd.read_file(os.path.join(
            self.input_destination_2, 'commercial/commercial.shp'))
        dfe = gpd.read_file(os.path.join(
            self.input_destination_2, 'educational/educational.shp'))
        dfi = gpd.read_file(os.path.join(
            self.input_destination_2, 'industrial/industrial.shp'))
        dfr = gpd.read_file(os.path.join(
            self.input_destination_2, 'residential/residential.shp'))
        self.area_a, self.area_c = dfa['area'].sum(), dfc['area'].sum()
        self.area_e, self.area_i = dfe['area'].sum(), dfi['area'].sum()
        self.area_r = dfr['area'].sum()
        self.area_a_pv, self.area_c_pv = self.area_a*0.267, self.area_c*0.267
        self.area_e_pv, self.area_i_pv = self.area_e*0.267, self.area_i*0.267
        self.area_r_pv = self.area_r*0.578
        # calculate aggregate peak power at rooftop
        total_roof_to_area = self.area_a_pv+self.area_c_pv + \
            self.area_e_pv+self.area_i_pv+self.area_r_pv
        self.peak_power_agg = (
            self.peak_power*(total_roof_to_area/self.module_area)) / self.norm  # KW
        print('Info: Maximum installed pv capacity = {} KW'.format(
            str(self.peak_power_agg)))
        # KW TODO: proper calculate max. installed cpacity for wind
        max_installed_wind_power = 305000
        maximum_capacity = [max_installed_wind_power, self.peak_power_agg]
        with open(os.path.join(self.output_destination, 'mc'), 'wb') as mc:
            pickle.dump(maximum_capacity, mc)
        logging.info("Calculate maximum installed PV capacity at roof top.")

    def agricultural_load(self):
        """Simulate Agricultural building type electricity demand and PV feedin supply."""
        load_a = self.time_stamp
        load_a['Load[kWh]'] = self.AL*self.area_a*self.x_a
        load_a['PV[kWh]'] = (self.pv_data*self.peak_power *
                             (self.area_a_pv/self.module_area)) / self.norm  # kWh
        load_a.to_csv(os.path.join(self.output_destination, 'a_load.csv'))
        logging.info(
            "Calculate electricty demand and supply for agricultural buildings")
        print('Info: Agricultural building load and pv supply simulation.')

    def commercial_energy_req(self):
        """Simulate Com.  building type electricity demand and PV feedin supply."""
        load_c = self.time_stamp
        load_c['Load[kWh]'] = self.CL*self.area_c*self.x_c
        load_c['PV[kWh]'] = (self.pv_data*self.peak_power *
                             (self.area_c_pv/self.module_area))/self.norm
        load_c.to_csv(os.path.join(self.output_destination, 'c_load.csv'))
        logging.info(
            "Calculate electricty demand and supply for commercial buildings")
        print('Info: Commercial building load and pv supply siml.')

    def eductaional_energy_req(self):
        """Simulate edu. quarter hourly Energy Requirments REs."""
        load_e = self.time_stamp
        load_e['Load[kWh]'] = self.EL*self.area_e*self.x_e
        load_e['PV[kWh]'] = (self.pv_data*self.peak_power *
                             (self.area_e_pv/self.module_area))/self.norm
        load_e.to_csv(os.path.join(self.output_destination, 'e_load.csv'))
        logging.info(
            "Calculate electricty demand and supply for educational buildings")
        print('Info: Eductaional building load and pv supply siml.')

    def industrial_energy_req(self):
        """Simulate Ind. quarter hourly Energy Requirments REs."""
        load_i = self.time_stamp
        load_i['Load[kWh]'] = self.IL*self.area_i*self.x_i
        load_i['PV[kWh]'] = (self.pv_data*self.peak_power *
                             (self.area_i_pv/self.module_area))/self.norm
        load_i.to_csv(os.path.join(self.output_destination, 'i_load.csv'))
        logging.info(
            "Calculate electricty demand and supply for industrial buildings")
        print('Info: Industrial building load and pv supply siml.')

    def residential_energy_req(self):
        """Simulate Res. quarter hourly Energy Requirments REs."""
        load_r = self.time_stamp
        load_r['Load[kWh]'] = self.RL*self.area_r*self.x_r
        load_r['PV[kWh]'] = (self.pv_data*self.peak_power *
                             (self.area_r_pv/self.module_area))/self.norm
        load_r.to_csv(os.path.join(self.output_destination, 'r_load.csv'))
        logging.info(
            "Calculate electricty demand and supply for residential buildings")
        print('Info: Residential building load and pv supply siml.')

    def highway_energy_req(self):
        """Simulate Urban Streetlightning. quarter hourly."""
        street_data = gpd.read_file(os.path.join(
            self.input_destination_2, 'highway/highway.shp'))
        area_sl = street_data['area'].sum()
        self.no_building = len(street_data)
        load_sl = self.time_stamp
        load_sl['Load[kWh]'] = self.SL*area_sl*self.x_sl
        load_sl.to_csv(os.path.join(self.output_destination, 'sl_load.csv'))
        logging.info(
            "Calculate electricty demand for street lights")
        print('Info: Streetlightning load siml.')

    def aggregate_demand_supply(self):
        """Aggrgate all simulated PV power generation."""
        print('Info: Aggregate simulated PV power generation and electricity demand.')

        demand_supply_agri = pd.read_csv(
            os.path.join(self.output_destination, 'a_load.csv'))
        demand_supply_comm = pd.read_csv(
            os.path.join(self.output_destination, 'c_load.csv'))
        demand_supply_educ = pd.read_csv(
            os.path.join(self.output_destination, 'e_load.csv'))
        demand_supply_indu = pd.read_csv(
            os.path.join(self.output_destination, 'i_load.csv'))
        demand_supply_resi = pd.read_csv(
            os.path.join(self.output_destination, 'r_load.csv'))
        demand_street_light = pd.read_csv(
            os.path.join(self.output_destination, 'sl_load.csv'))

        # aggregate all pv supply for the different building types
        agg_pv = (demand_supply_agri['PV[kWh]'] + demand_supply_comm['PV[kWh]'] +
                  demand_supply_educ['PV[kWh]'] + demand_supply_indu['PV[kWh]'] +
                  demand_supply_resi['PV[kWh]'])/self.norm  # to MWh
        agg_pv = np.array(agg_pv)
        agg_pv = pd.DataFrame(agg_pv, columns=["PV[MWh]"])

        # aggregate all electricity demand for the different building types (MWh)
        agg_load = (demand_supply_agri['Load[kWh]'] + demand_supply_comm['Load[kWh]'] +
                    demand_supply_educ['Load[kWh]'] + demand_supply_indu['Load[kWh]'] +
                    demand_supply_resi['Load[kWh]'] + demand_street_light['Load[kWh]'])/self.norm

        # prepare demand and supply data for optimization
        agg_load = np.array(agg_load)
        agg_load = pd.DataFrame(agg_load, columns=["Load[MWh]"])
        agg_load['PV[MWh]'] = agg_pv["PV[MWh]"]
        agg_load['demand_el'] = agg_load["Load[MWh]"].values * \
            1000  # Kilo-watts hour
        agg_load['wind'] = self.wind_power_data  # normalized wind feedin data
        agg_load['pv'] = self.pv_data  # normalized pv feedin data
        agg_load['time'] = self.time_stamp['time']
        agg_load = agg_load.set_index('time')

        agg_load.loc[:, ["Load[MWh]", 'PV[MWh]']].to_csv(os.path.join(
            self.output_destination, 'aggregated-demand-supply.csv'))
        agg_load.loc[:, ["demand_el", 'wind', 'pv']].to_csv(
            os.path.join(self.output_destination, 'optimization-commodities.csv'))
        logging.info(
            "Calculate aggregated electricty demand and supplies")

    def plot_quarter_load_energy(self):
        """Simulate quarter load and plot Urban Energy Requirments REs."""
        print('Info: Plot Urban Energy Requirments.')
        fig_size = (14, 8)
        sim_df = pd.read_csv(os.path.join(
            self.output_destination, 'aggregated-demand-supply.csv'))
        sim_df['PV[MWh]'].plot(style='r', figsize=fig_size, grid=True)
        sim_df['Load[MWh]'].plot(style='g', figsize=fig_size, grid=True)
        plt.xlabel('Time')
        plt.ylabel('MW')
        plt.title('Aggregated Energy Requirments in Oldenburg')
        plt.legend(['Simulated PV', 'Simulated load'], loc='upper left')
        plt.savefig(self.output_destination2+"Energy_Requirments.png", dpi=300)
        plt.show()
        logging.info("Generate demand and supply plot")

    def flexigis_urban_simulation(self):
        """Simulate urban Energy requirements for target location."""
        self.feedin_data()
        self.standard_load_profiles()
        self.electricity_usage_index()
        self.roofTop_pv_capacity()
        self.agricultural_load()
        self.commercial_energy_req()
        self.eductaional_energy_req()
        self.industrial_energy_req()
        self.residential_energy_req()
        self.highway_energy_req()
        self.aggregate_demand_supply()
        self.plot_quarter_load_energy()
        logging.info("Electricity Demand and Supply simulation Done!")

if __name__ == "__main__":
    flexigis_energy = UrbanEnergyRequirement()
    flexigis_energy.flexigis_urban_simulation()


#! /usr/bin/env python

########################################################################
#                                                                      #
# Resums the non-global logarithms, needs ngl_resum.py                 #
#                                                                      #
# If using ngl_resum, please cite                                      #
#               doi:10.1007/JHEP09(2020)029                            #
#               https://inspirehep.net/literature/1798660              #
#                                                                      #
########################################################################

__author__ = 'Marcel Balsiger'
__email__ = 'marcel.balsiger@hotmail.com'
__date__ = 'October 19, 2020'

import time
import numpy as np
import argparse
import pylhe
import ngl_resum as ngl

parser = argparse.ArgumentParser(description='This code shows how to '\
        'use ngl_resum in combination with LHE-files, considering '\
        'top-pair production. First, each event gets tested whether '\
        'it fulfills the conditions of Table 1 from the ATLAS paper '\
        'arXiv:1203.5015 [hep-ex] and then showers the dipoles with '\
        'the outside region defined by the symmetric rapidity gap '\
        'from -y to y with areas around the bottom quarks cut away. '\
        'Similar code was used to resum the non-global logarithms in '\
        'Section 5 of arXiv:2006.00014')
parser.add_argument('-f','--file', help='lhe event file to shower',\
        default=None,required=True)
parser.add_argument('-y','--ymax', help='ymax of outside region', \
        default=0.8, type=float)
parser.add_argument('-n','--nsh', help='number of showers per dipole', \
        default=100, type=int)
parser.add_argument('-t','--tmax', help='maximal shower time tmax', \
        default=0.1, type=float)
parser.add_argument('-m','--nbins', help='number of bins in hists', \
        default=100, type=int)
parser.add_argument('-c','--cutoff', help='cutoff of shower', \
        default=5, type=float)
parser.add_argument('-s','--seed', help='random seed', \
        default=None, type=int)
parser.add_argument('-b','--break', help='stop after so many events', \
        default=100000, type=int)
args = vars(parser.parse_args())
eventFile=args['file']

if not(args['seed'] is None) : np.random.seed(args['seed'])

showerCutoff=float(args['cutoff'])
nbins=int(args['nbins'])
tmax=float(args['tmax'])
nsh=int(args['nsh'])

def _outside(self,v):
    jetaxis1=self.event.outgoingBottom[0]/self.event.outgoingBottom[0].e
    jetaxis2=self.event.outgoingBottom[1]/self.event.outgoingBottom[1].e
    jetRadius=0.4
    rapRangeMax=float(args['ymax'])
    rapRangeMin=0.0
    return (v.R2(jetaxis1)>jetRadius**2) and \
           (v.R2(jetaxis2)>jetRadius**2) and \
           (abs(v.rap)<rapRangeMax) and (abs(v.rap)>=rapRangeMin)

def validEvent(ev): # ev is the ngl.Event we want to test

    # check whether we have the necessary particles
    if ev.intermediateTop == None : return False
    if ev.outgoingBottom == None : return False
    if (ev.outgoingElectron == None) and (ev.outgoingMuon == None): \
            return False
    if len(ev.intermediateTop) != 2 :  return False
    if len(ev.outgoingBottom) != 2 : return False

    momentaLeptonsOut=[]
    momentaNeutrinoOut=[]

    electronmuonevent=True
    if not ev.outgoingElectron==None:
        for i in ev.outgoingElectron:
            momentaLeptonsOut.append(i)
            # checks on electron(s)
            if i.eT< 25: return False
            if abs(i.rap)>2.47: return False
        for i in ev.outgoingENeutrino:
            momentaNeutrinoOut.append(i)
    else:
        electronmuonevent=False

    if not ev.outgoingMuon==None:
        for i in ev.outgoingMuon:
            momentaLeptonsOut.append(i)
            # checks on muon(s)
            if i.pT< 20: return False
            if abs(i.rap)>2.5: return False
        for i in ev.outgoingMNeutrino:
            momentaNeutrinoOut.append(i)
    else:
        electronmuonevent=False

    # check number of leptons ans neutrinos
    if len(momentaLeptonsOut) != 2 : return False
    if len(momentaNeutrinoOut) != 2 : return False

    dileptonmass=np.sqrt((momentaLeptonsOut[0]+momentaLeptonsOut[1])*\
                            (momentaLeptonsOut[0]+momentaLeptonsOut[1]))
    missingMomentum=(momentaNeutrinoOut[0]+momentaNeutrinoOut[1])

    if not electronmuonevent:
        # checks on "missing momenta" (neutrinos) and dilepton mass
        if missingMomentum.eT<40 : return False
        if (dileptonmass<15 or abs(dileptonmass-91)<10) : return False
    else:
        # check on visible transverse momentum
        if (momentaLeptonsOut[0].pT+momentaLeptonsOut[1].pT+\
            ev.outgoingBottom[0].pT+ev.outgoingBottom[1].pT)<130:
                return False

    # checks on bottom quarks
    for i in ev.outgoingBottom:
        if i.pT<25: return False
        if abs(i.rap)>2.4: return False
        for j in momentaLeptonsOut:
                    if i.R2(j)<0.4**2: return False

    return True # only gets reached, if no check failed.

evtFile = pylhe.readLHE(eventFile)

fullResultLL=ngl.Hist(nbins,tmax,errorHistCalc=True)
fullNGL1Loop=0.
fullNGL1LoopSq=0.
fullNGL2Loop=0.
fullNGL2LoopSq=0.

eventWeight=0.

numberEvents=0
numberValidEvents=0

timeStart = time.time()

for event in evtFile:
    numberEvents+=1

    ev=ngl.Event(eventFromFile=event,productionDipoles='intermediate',\
                    decayDipoles=False)

    if not eventWeight > 0:
        eventWeight=ev.weight
    if not eventWeight==ev.weight:
        print("Warning: events not of equal weight!")

    if validEvent(ev):
        numberValidEvents+=1

        outsideRegion=ngl.OutsideRegion(ev)
        outsideRegion.outside = _outside.__get__(outsideRegion,\
            ngl.OutsideRegion)
        shower=ngl.Shower(ev,outsideRegion,nsh,nbins,tmax,showerCutoff)
        shower.shower()
        fullResultLL+=shower.resLL
        fullNGL1Loop+=shower.ngl1Loop
        fullNGL1LoopSq+=shower.ngl1LoopSq
        fullNGL2Loop+=shower.ngl2Loop
        fullNGL2LoopSq+=shower.ngl2LoopSq

    if numberEvents >= int(args['break']):break

print('runtime=', time.time()-timeStart,' sec')
print("of ", numberEvents," events, ", numberValidEvents," were valid.")
print("Weight of each event:", eventWeight)
print('\n\n'  )

print('*************************************')
print('*  t       LL(t)          dS(t)     * ')
print('*************************************\n')
print('*** Binned Result ***\n\n')

for i in range(0,fullResultLL.nbins):
    print( round(fullResultLL.centerBinValue[i],4),' ', \
                fullResultLL.entries[i]/numberValidEvents,' ', \
                np.sqrt(fullResultLL.squaredError[i])/numberValidEvents)

print('\n'  )
snlo=fullNGL1Loop/numberValidEvents
snloError=np.sqrt((fullNGL1LoopSq/numberValidEvents-\
                        (fullNGL1Loop/numberValidEvents)**2)\
                        /(nsh*numberValidEvents))
print('snlo=',snlo)
print('snloError=',snloError)

print('\n')

snnlo=fullNGL2Loop/numberValidEvents+0.5*snlo**2
#Error(snnlo)=|d(snnlo)/d(fullNGL2Loop)*Error(fullNGL2Loop)|
#               + |d(snnlo)/d(snlo)*Error(snlo)|
snnloError=abs(np.sqrt((fullNGL2LoopSq/numberValidEvents-\
                (fullNGL2Loop/numberValidEvents)**2)/\
                (nsh*numberValidEvents)))\
                +abs(snlo*snloError)
print('snnlo=',snnlo)
print('snnloError=',snnloError)
print('\n')


import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Layer
from tensorflow.keras import Sequential
import tensorflow.keras.layers as nn
from tensorflow import einsum
from tensorflow.keras.preprocessing.sequence import pad_sequences

from einops import rearrange, repeat
from einops.layers.tensorflow import Rearrange

import numpy as np

def exists(val):
    return val is not None

def pair(t):
    return t if isinstance(t, tuple) else (t, t)

# adaptive token sampling functions and classes

def log(t, eps = 1e-6):
    return tf.math.log(t + eps)

def sample_gumbel(shape, dtype, eps = 1e-6):
    u = tf.random.uniform(shape, dtype=dtype)
    return -log(-log(u, eps), eps)

def torch_gather(x, indices, gather_axis):
    # if pytorch gather indices are
    # [[[0, 10, 20], [0, 10, 20], [0, 10, 20]],
    #  [[0, 10, 20], [0, 10, 20], [0, 10, 20]]]
    # tf nd_gather needs to be
    # [[0,0,0], [0,0,10], [0,0,20], [0,1,0], [0,1,10], [0,1,20], [0,2,0], [0,2,10], [0,2,20],
    #  [1,0,0], [1,0,10], [1,0,20], [1,1,0], [1,1,10], [1,1,20], [1,2,0], [1,2,10], [1,2,20]]

    indices = tf.cast(indices, tf.int64)
    # create a tensor containing indices of each element
    all_indices = tf.where(tf.fill(indices.shape, True))
    gather_locations = tf.reshape(indices, [indices.shape.num_elements()])

    # splice in our pytorch style index at the correct axis
    gather_indices = []
    for axis in range(len(indices.shape)):
        if axis == gather_axis:
            gather_indices.append(gather_locations)
        else:
            gather_indices.append(all_indices[:, axis])

    gather_indices = tf.stack(gather_indices, axis=-1)
    gathered = tf.gather_nd(x, gather_indices)
    reshaped = tf.reshape(gathered, indices.shape)
    return reshaped

def batched_index_select(values, indices, dim = 1):
    value_dims = values.shape[(dim + 1):]
    values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
    indices = indices[(..., *((None,) * len(value_dims)))]
    indices = tf.tile(indices, multiples=[1] * len(indices_shape) + [*value_dims])
    value_expand_len = len(indices_shape) - (dim + 1)
    values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]

    value_expand_shape = [-1] * len(values.shape)
    expand_slice = slice(dim, (dim + value_expand_len))
    value_expand_shape[expand_slice] = indices.shape[expand_slice]
    dim += value_expand_len

    values = torch_gather(values, indices, dim)
    return values

class AdaptiveTokenSampling(Layer):
    def __init__(self, output_num_tokens, eps=1e-6):
        super(AdaptiveTokenSampling, self).__init__()
        self.eps = eps
        self.output_num_tokens = output_num_tokens

    def call(self, attn, value=None, mask=None, training=True):
        heads, output_num_tokens, eps, dtype = attn.shape[1], self.output_num_tokens, self.eps, attn.dtype

        # first get the attention values for CLS token to all other tokens
        cls_attn = attn[..., 0, 1:]

        # calculate the norms of the values, for weighting the scores, as described in the paper
        value_norms = tf.norm(value[..., 1:, :], axis=-1)

        # weigh the attention scores by the norm of the values, sum across all heads
        cls_attn = einsum('b h n, b h n -> b n', cls_attn, value_norms)

        # normalize to 1
        normed_cls_attn = cls_attn / (tf.reduce_sum(cls_attn, axis=-1, keepdims=True) + eps)

        # instead of using inverse transform sampling, going to invert the softmax and use gumbel-max sampling instead
        pseudo_logits = log(normed_cls_attn)

        # mask out pseudo logits for gumbel-max sampling
        mask_without_cls = mask[:, 1:]
        mask_value = -np.finfo(attn.dtype.as_numpy_dtype).max / 2
        pseudo_logits = tf.where(~mask_without_cls, mask_value, pseudo_logits)

        # expand k times, k being the adaptive sampling number
        pseudo_logits = repeat(pseudo_logits, 'b n -> b k n', k=output_num_tokens)
        pseudo_logits = pseudo_logits + sample_gumbel(pseudo_logits.shape, dtype=dtype)

        # gumble-max and add one to reserve 0 for padding / mask
        sampled_token_ids = tf.argmax(pseudo_logits, axis=-1) + 1

        # calculate unique using torch.unique and then pad the sequence from the right
        unique_sampled_token_ids_list = []
        for t in tf.unstack(sampled_token_ids):
            t = tf.cast(t, tf.int32)
            t, _ = tf.unique(t)
            x = tf.sort(t)
            unique_sampled_token_ids_list.append(x)

        unique_sampled_token_ids = pad_sequences(unique_sampled_token_ids_list)

        # calculate the new mask, based on the padding
        new_mask = unique_sampled_token_ids != 0

        # CLS token never gets masked out (gets a value of True)
        new_mask = tf.pad(new_mask, paddings=[[0, 0], [1, 0]], constant_values=True)

        # prepend a 0 token id to keep the CLS attention scores
        unique_sampled_token_ids = tf.pad(unique_sampled_token_ids, paddings=[[0, 0], [1, 0]])
        expanded_unique_sampled_token_ids = repeat(unique_sampled_token_ids, 'b n -> b h n', h=heads)

        # gather the new attention scores
        new_attn = batched_index_select(attn, expanded_unique_sampled_token_ids, dim=2)

        # return the sampled attention scores, new mask (denoting padding), as well as the sampled token indices (for the residual)
        return new_attn, new_mask, unique_sampled_token_ids

def gelu(x, approximate=False):
    if approximate:
        coeff = tf.cast(0.044715, x.dtype)
        return 0.5 * x * (1.0 + tf.tanh(0.7978845608028654 * (x + coeff * tf.pow(x, 3))))
    else:
        return 0.5 * x * (1.0 + tf.math.erf(x / tf.cast(1.4142135623730951, x.dtype)))

class GELU(Layer):
    def __init__(self, approximate=False):
        super(GELU, self).__init__()
        self.approximate = approximate

    def call(self, x, training=True):
        return gelu(x, self.approximate)

class PreNorm(Layer):
    def __init__(self, fn):
        super(PreNorm, self).__init__()

        self.norm = nn.LayerNormalization()
        self.fn = fn

    def call(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class MLP(Layer):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super(MLP, self).__init__()
        self.net = Sequential([
            nn.Dense(units=hidden_dim),
            GELU(),
            nn.Dropout(rate=dropout),
            nn.Dense(units=dim),
            nn.Dropout(rate=dropout)
        ])

    def call(self, x, training=True):
        return self.net(x, training=training)

class Attention(Layer):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, output_num_tokens=None):
        super(Attention, self).__init__()
        inner_dim = dim_head * heads
        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax()
        self.to_qkv = nn.Dense(units=inner_dim * 3, use_bias=False)

        self.output_num_tokens = output_num_tokens
        self.ats = AdaptiveTokenSampling(output_num_tokens) if exists(output_num_tokens) else None

        self.to_out = Sequential([
            nn.Dense(units=dim),
            nn.Dropout(rate=dropout)
        ])

    def call(self, x, mask=None, training=True):
        num_tokens = x.shape[1]

        qkv = self.to_qkv(x)
        qkv = tf.split(qkv, num_or_size_splits=3, axis=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d)-> b h n d', h=self.heads), qkv)

        dots = tf.matmul(q, tf.transpose(k, perm=[0, 1, 3, 2])) * self.scale

        if exists(mask):
            mask_f = tf.cast(mask, tf.float32)
            dots_mask = rearrange(mask_f, 'b i -> b 1 i 1') * rearrange(mask_f, 'b j -> b 1 1 j')
            dots_mask = tf.cast(dots_mask, tf.bool)
            mask_value = -np.finfo(dots.dtype.as_numpy_dtype).max
            dots = tf.where(~dots_mask, mask_value, dots)

        attn = self.attend(dots)

        sampled_token_ids = None

        # if adaptive token sampling is enabled
        # and number of tokens is greater than the number of output tokens
        if exists(self.output_num_tokens) and (num_tokens - 1) > self.output_num_tokens:
            attn, mask, sampled_token_ids = self.ats(attn, v, mask=mask)

        out = tf.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out = self.to_out(out, training=training)

        return out, mask, sampled_token_ids

class Transformer(Layer):
    def __init__(self, dim, depth, max_tokens_per_depth, heads, dim_head, mlp_dim, dropout=0.0):
        super(Transformer, self).__init__()
        assert len(max_tokens_per_depth) == depth, 'max_tokens_per_depth must be a tuple of length that is equal to the depth of the transformer'
        assert sorted(max_tokens_per_depth, reverse=True) == list(max_tokens_per_depth), 'max_tokens_per_depth must be in decreasing order'
        assert min(max_tokens_per_depth) > 0, 'max_tokens_per_depth must have at least 1 token at any layer'

        self.layers = []
        for _, output_num_tokens in zip(range(depth), max_tokens_per_depth):
            self.layers.append([
                PreNorm(Attention(dim, output_num_tokens=output_num_tokens, heads=heads, dim_head=dim_head, dropout=dropout)),
                PreNorm(MLP(dim, mlp_dim, dropout=dropout))
            ])

    def call(self, x, training=True):
        b, n = x.shape[:2]

        # use mask to keep track of the paddings when sampling tokens
        # as the duplicates (when sampling) are just removed, as mentioned in the paper
        mask = tf.ones([b, n], dtype=tf.bool)

        token_ids = tf.range(n)
        token_ids = repeat(token_ids, 'n -> b n', b = b)

        for attn, ff in self.layers:
            attn_out, mask, sampled_token_ids = attn(x, mask=mask, training=training)

            # when token sampling, one needs to then gather the residual tokens with the sampled token ids
            if exists(sampled_token_ids):
                x = batched_index_select(x, sampled_token_ids, dim=1)
                token_ids = batched_index_select(token_ids, sampled_token_ids, dim=1)

            x = x + attn_out

            x = ff(x, training=training) + x

        return x, token_ids

class ViT(Model):
    def __init__(self,
                 image_size,
                 patch_size,
                 num_classes,
                 dim,
                 depth,
                 max_tokens_per_depth,
                 heads,
                 mlp_dim,
                 dim_head=64,
                 dropout=0.0,
                 emb_dropout=0.0
                 ):
        super(ViT, self).__init__()

        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(patch_size)

        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'

        num_patches = (image_height // patch_height) * (image_width // patch_width)

        self.patch_embedding = Sequential([
            Rearrange('b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),
            nn.Dense(units=dim)
        ])

        self.pos_embedding = tf.Variable(initial_value=tf.random.normal([1, num_patches + 1, dim]))
        self.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]))
        self.dropout = nn.Dropout(rate=emb_dropout)

        self.transformer = Transformer(dim, depth, max_tokens_per_depth, heads, dim_head, mlp_dim, dropout)

        self.mlp_head = Sequential([
            nn.LayerNormalization(),
            nn.Dense(units=num_classes)
        ])

    def call(self, img, return_sampled_token_ids=False, training=True, **kwargs):
        x = self.patch_embedding(img)
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
        x = tf.concat([cls_tokens, x], axis=1)
        x += self.pos_embedding[:, :(n + 1)]
        x = self.dropout(x, training=training)

        x, token_ids = self.transformer(x, training=training)

        logits = self.mlp_head(x[:, 0])

        if return_sampled_token_ids:
            # remove CLS token and decrement by 1 to make -1 the padding
            token_ids = token_ids[:, 1:] - 1
            return logits, token_ids

        return logits

v = ViT(
    image_size = 256,
    patch_size = 16,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    max_tokens_per_depth = (256, 128, 64, 32, 16, 8), # a tuple that denotes the maximum number of tokens that any given layer should have. if the layer has greater than this amount, it will undergo adaptive token sampling
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

img = tf.random.normal(shape=[4, 256, 256, 3])
preds = v(img) # (1, 1000)
print(preds.shape)


"""
kaldo
Anharmonic Lattice Dynamics
"""
import numpy as np
from sparse import COO
import pandas as pd
import ase.units as units
from kaldo.helpers.tools import count_rows
from ase import Atoms
import re
from kaldo.helpers.logger import get_logger
logging = get_logger()

def import_from_files(replicated_atoms, dynmat_file=None, third_file=None, supercell=(1, 1, 1),
                      third_energy_threshold=0.):
    # TODO: split this method into two pieces
    n_replicas = np.prod(supercell)
    n_total_atoms = replicated_atoms.positions.shape[0]
    n_unit_atoms = int(n_total_atoms / n_replicas)
    unit_symbols = []
    unit_positions = []
    for i in range(n_unit_atoms):
        unit_symbols.append(replicated_atoms.get_chemical_symbols()[i])
        unit_positions.append(replicated_atoms.positions[i])
    unit_cell = replicated_atoms.cell / supercell

    atoms = Atoms(unit_symbols,
                  positions=unit_positions,
                  cell=unit_cell,
                  pbc=[1, 1, 1])

    second_order = None
    third_order = None

    if dynmat_file:
        logging.info('Reading dynamical matrix')
        second_dl = import_second(atoms, replicas=supercell, filename=dynmat_file)
        second_order = second_dl

    if third_file:
        try:
            logging.info('Reading sparse third order')
            third_dl = import_sparse_third(atoms=atoms,
                                              supercell=supercell,
                                              filename=third_file,
                                              third_energy_threshold=third_energy_threshold)

        except UnicodeDecodeError:
            if third_energy_threshold != 0:
                raise ValueError('Third threshold not supported for dense third')
            logging.info('Reading dense third order')
            third_dl = import_dense_third(atoms, supercell=supercell, filename=third_file)
            logging.info('Third order matrix stored.')
        third_dl = third_dl[:n_unit_atoms]
        third_shape = (
            n_unit_atoms * 3, n_replicas * n_unit_atoms * 3, n_replicas * n_unit_atoms * 3)
        third_dl = third_dl.reshape(third_shape)
        third_order = third_dl

    return second_order, third_order

def import_second(atoms, replicas=(1, 1, 1), filename='Dyn.form'):
    replicas = np.array(replicas)
    n_unit_cell = atoms.positions.shape[0]
    dyn_mat = import_dynamical_matrix(n_unit_cell, replicas, filename)
    mass = np.sqrt (atoms.get_masses ())
    dyn_mat = dyn_mat * mass[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis, np.newaxis]
    dyn_mat = dyn_mat * mass[np.newaxis, np.newaxis, np.newaxis, np.newaxis, :, np.newaxis]
    return dyn_mat

def import_dynamical_matrix(n_atoms, supercell=(1, 1, 1), filename='Dyn.form'):
    supercell = np.array(supercell)
    dynamical_matrix_frame = pd.read_csv(filename, header=None, delim_whitespace=True)
    dynamical_matrix = dynamical_matrix_frame.values
    n_replicas = np.prod(supercell)
    if dynamical_matrix.size == n_replicas * (n_atoms * 3) ** 2:
        dynamical_matrix = dynamical_matrix.reshape((n_atoms, 3, n_replicas, n_atoms, 3))
    elif dynamical_matrix.size == (n_replicas * n_atoms * 3) ** 2:
        dynamical_matrix = dynamical_matrix.reshape((n_replicas, n_atoms, 3, n_replicas, n_atoms, 3))[0]
    elif dynamical_matrix.size == (n_atoms * 3) ** 2:
        dynamical_matrix = dynamical_matrix.reshape((n_atoms, 3, 1, n_atoms, 3))
    else:
        logging.error('Impossible to read calculate_dynmat with size ' + str(dynamical_matrix.size))
    tenjovermoltoev = 10 * units.J / units.mol
    return dynamical_matrix * tenjovermoltoev

def import_sparse_third(atoms, supercell=(1, 1, 1), filename='THIRD', third_energy_threshold=0.):
    supercell = np.array(supercell)
    n_replicas = np.prod(supercell)
    n_atoms = atoms.get_positions().shape[0]
    n_replicated_atoms = n_atoms * n_replicas
    n_rows = count_rows(filename)
    array_size = min(n_rows * 3, n_atoms * 3 * (n_replicated_atoms * 3) ** 2)
    coords = np.zeros((array_size, 6), dtype=np.int16)
    values = np.zeros((array_size))
    index_in_unit_cell = 0
    tenjovermoltoev = 10 * units.J / units.mol
    with open(filename) as f:
        for i, line in enumerate(f):
            l_split = re.split('\s+', line.strip())
            coords_to_write = np.array(l_split[0:-3], dtype=int) - 1
            values_to_write = np.array(l_split[-3:], dtype=np.float)
            #TODO: add 'if' third_energy_threshold before calculating the mask
            mask_to_write = np.abs(values_to_write) > third_energy_threshold

            if mask_to_write.any() and coords_to_write[0] < n_atoms:
                for alpha in np.arange(3)[mask_to_write]:
                    coords[index_in_unit_cell, :-1] = coords_to_write[np.newaxis, :]
                    coords[index_in_unit_cell, -1] = alpha
                    values[index_in_unit_cell] = values_to_write[alpha] * tenjovermoltoev
                    index_in_unit_cell = index_in_unit_cell + 1
            if i % 1000000 == 0:
                logging.info('reading third order: ' + str(np.round(i / n_rows, 2) * 100) + '%')
    logging.info('read ' + str(3 * i) + ' interactions')
    coords = coords[:index_in_unit_cell].T
    values = values[:index_in_unit_cell]
    sparse_third = COO (coords, values, shape=(n_atoms, 3, n_replicated_atoms, 3, n_replicated_atoms, 3))
    return sparse_third

def import_dense_third(atoms, supercell, filename, is_reduced=True):
    supercell = np.array(supercell)
    n_replicas = np.prod(supercell)
    n_atoms = atoms.get_positions().shape[0]
    if is_reduced:
        total_rows = (n_atoms *  3) * (n_atoms * n_replicas * 3) ** 2
        third = np.fromfile(filename, dtype=np.float, count=total_rows)
        third = third.reshape((n_atoms, 3, n_atoms * n_replicas, 3, n_atoms * n_replicas, 3))
    else:
        total_rows = (n_atoms * n_replicas * 3) ** 3
        third = np.fromfile(filename, dtype=np.float, count=total_rows)
        third = third.reshape((n_atoms * n_replicas, 3, n_atoms * n_replicas, 3, n_atoms * n_replicas, 3))
    return third


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

class periodic_box_condition_information:
    """periodic_box_condition_information"""

    def __init__(self, controller, box_length):
        CONSTANT_UINT_MAX_FLOAT = 4294967296.0
        self.crd_to_uint_crd_cof = np.array([CONSTANT_UINT_MAX_FLOAT / box_length[0],
                                             CONSTANT_UINT_MAX_FLOAT / box_length[1],
                                             CONSTANT_UINT_MAX_FLOAT / box_length[2]])
        self.quarter_crd_to_uint_crd_cof = 0.25 * self.crd_to_uint_crd_cof
        self.uint_dr_to_dr_cof = 1.0 / self.crd_to_uint_crd_cof

class system_information:
    """system_information"""

    def __init__(self, controller, md_info):
        CONSTANT_PRES_CONVERTION_INVERSE = 0.00001439506089041446
        self.md_info = md_info
        self.box_length = self.md_info.box_length
        self.steps = 0
        self.step_limit = 1000 if "step_limit" not in controller.Command_Set else int(
            controller.Command_Set["step_limit"])
        self.target_temperature = 300.0 if "target_temperature" not in controller.Command_Set else float(
            controller.Command_Set["target_temperature"])
        if md_info.mode == 2 and "target_pressure" in controller.Command_Set:
            self.target_pressure = float(controller.Command_Set["target_pressure"])
        else:
            self.target_pressure = 1
        self.target_pressure *= CONSTANT_PRES_CONVERTION_INVERSE
        self.d_virial = 0
        self.d_pressure = 0
        self.d_temperature = 0
        self.d_potential = 0
        self.d_sum_of_atom_ek = 0
        self.freedom = 3 * md_info.atom_numbers

class non_bond_information:
    """system_information"""

    def __init__(self, controller, md_info):
        self.md_info = md_info
        self.skin = 2.0 if "skin" not in controller.Command_Set else float(controller.Command_Set["skin"])
        print("    skin set to %.2f Angstram" % (self.skin))
        self.cutoff = 10.0 if "cutoff" not in controller.Command_Set else float(controller.Command_Set["cutoff"])
        self.atom_numbers = self.md_info.atom_numbers
        self.excluded_atom_numbers = 0
        self.h_excluded_list_start = []
        self.h_excluded_numbers = []
        self.h_excluded_list = []
        if controller.amber_parm is not None:
            file_path = controller.amber_parm
            self.read_information_from_amberfile(file_path)
        else:
            self.read_exclude_file(controller)

    def read_exclude_file(self, controller):
        """read_exclude_file"""
        if "exclude_in_file" in controller.Command_Set:
            print("    Start reading excluded list:")
            path = controller.Command_Set["exclude_in_file"]
            file = open(path, 'r')
            context = file.readlines()
            atom_numbers, self.excluded_atom_numbers = list(map(int, context[0].strip().split()))
            if self.md_info.atom_numbers > 0 and (atom_numbers != self.md_info.atom_numbers):
                print("        Error: atom_numbers is not equal: ", atom_numbers, self.md_info.atom_numbers)
                exit(1)
            else:
                self.md_info.atom_numbers = atom_numbers
            count = 0
            for idx, val in enumerate(context):
                if idx > 0:
                    el = list(map(int, val.strip().split()))
                    if el[0] == 1 and -1 in el:
                        self.h_excluded_numbers.append(0)
                    else:
                        self.h_excluded_numbers.append(el[0])
                    self.h_excluded_list_start.append(count)
                    if el:
                        self.h_excluded_list.extend(el[1:])
                        count += el[0]
            print("    End reading excluded list")
            file.close()
        else:
            print("    Set all atom exclude no atoms as default")
            count = 0
            for i in range(self.md_info.atom_numbers):
                self.h_excluded_numbers[i] = 0
                self.h_excluded_list_start[i] = count
                for _ in range(self.h_excluded_numbers[i]):
                    self.h_excluded_list[count] = 0
                    count += 1
            print("    End reading charge")

    def read_information_from_amberfile(self, file_path):
        '''read amber file'''
        file = open(file_path, 'r')
        context = file.readlines()
        file.close()
        self.h_excluded_list_start = [0] * self.atom_numbers
        self.h_excluded_numbers = [0] * self.atom_numbers

        for idx, val in enumerate(context):
            if idx < len(context) - 1:
                if "%FLAG POINTERS" in val + context[idx + 1] and "%FORMAT(10I8)" in val + context[idx + 1]:
                    start_idx = idx + 2
                    count = 0
                    value = list(map(int, context[start_idx].strip().split()))
                    information = []
                    information.extend(value)
                    while count < 11:
                        start_idx += 1
                        value = list(map(int, context[start_idx].strip().split()))
                        information.extend(value)
                        count += len(value)
                    self.excluded_atom_numbers = information[10]
                    print("excluded atom numbers ", self.excluded_atom_numbers)
                    break
        for idx, val in enumerate(context):
            if "%FLAG NUMBER_EXCLUDED_ATOMS" in val:
                count = 0
                start_idx = idx
                information = []
                while count < self.atom_numbers:
                    start_idx += 1
                    if "%FORMAT" in context[start_idx]:
                        continue
                    else:
                        value = list(map(int, context[start_idx].strip().split()))
                        information.extend(value)
                        count += len(value)
                count = 0
                for i in range(self.atom_numbers):
                    self.h_excluded_numbers[i] = information[i]
                    self.h_excluded_list_start[i] = count
                    count += information[i]
                break

        total_count = sum(self.h_excluded_numbers)
        self.h_excluded_list = []
        for idx, val in enumerate(context):
            if "%FLAG EXCLUDED_ATOMS_LIST" in val:
                count = 0
                start_idx = idx
                information = []
                while count < total_count:
                    start_idx += 1
                    if "%FORMAT" in context[start_idx]:
                        continue
                    else:
                        value = list(map(int, context[start_idx].strip().split()))
                        information.extend(value)
                        count += len(value)

                count = 0
                for i in range(self.atom_numbers):
                    tmp_list = []
                    if self.h_excluded_numbers[i] == 1:
                        tmp_list.append(information[count] - 1)
                        if information[count] == 0:
                            self.h_excluded_numbers[i] = 0
                        count += 1
                    else:
                        for _ in range(self.h_excluded_numbers[i]):
                            tmp_list.append(information[count] - 1)

                            count += 1
                        tmp_list = sorted(tmp_list)
                    self.h_excluded_list.extend(tmp_list)
                break

class NVE_iteration:
    """NVE_iteration"""

    def __init__(self, controller, md_info):
        self.max_velocity = -1 if "nve_velocity_max" not in controller.Command_Set else float(
            controller.Command_Set["nve_velocity_max"])

class residue_information:
    """residue_information"""

    def __init__(self, controller, md_info):
        self.md_info = md_info
        self.residue_numbers = 0
        self.h_mass = []
        self.h_mass_inverse = []
        self.h_res_start = []
        self.h_res_end = []
        self.momentum = []
        self.center_of_mass = []
        self.sigma_of_res_ek = 0
        self.res_ek_energy = 0
        self.sigma_of_res_ek = 0
        self.is_initialized = 0
        print("    Start reading residue list:")
        if "residue_in_file" in controller.Command_Set:
            self.read_residule_file(controller)
        elif "amber_parm7" in controller.Command_Set:
            self.residue_numbers = self.md_info.residue_numbers
            self.h_res_start = md_info.h_res_start
            self.h_res_end = md_info.h_res_end
            self.is_initialized = 1
            self.read_res_mass()
        else:
            self.residue_numbers = md_info.atom_numbers
            self.h_res_start = list(range(self.residue_numbers))
            self.h_res_end = list(range(1, self.residue_numbers + 1))
            self.is_initialized = 1
        print("    End reading residue list")

    def read_res_mass(self):
        """ Read_AMBER_Parm7 """
        if self.md_info.h_mass:
            for i in range(self.residue_numbers):
                temp_mass = 0
                for j in range(self.h_res_start[i], self.h_res_end[i]):
                    temp_mass += self.md_info.h_mass[j]
                self.h_mass.append(temp_mass)
                if temp_mass == 0:
                    self.h_mass_inverse.append(0)
                else:
                    self.h_mass_inverse.append(1.0 / temp_mass)
        else:
            print("    Error: atom mass should be initialized before residue mass")
            exit(1)

    def read_residule_file(self, controller):
        """read_residule_file"""
        if "residue_in_file" in controller.Command_Set:
            path = controller.Command_Set["residue_in_file"]
            file = open(path, 'r')
            context = file.readlines()
            atom_numbers, self.residue_numbers = list(map(int, context[0].strip().split()))
            print("        residue_numbers is ", self.residue_numbers)
            # self.md_info.residue_numbers = self.residue_numbers
            if self.md_info.atom_numbers > 0 and (atom_numbers != self.md_info.atom_numbers):
                print("        Error: atom_numbers is not equal: ", atom_numbers, self.md_info.atom_numbers)
                exit(1)
            else:
                self.md_info.atom_numbers = atom_numbers
            print("        residue_numbers is ", self.residue_numbers)

            count = 0
            for idx, val in enumerate(context):
                if idx > 0:
                    self.h_res_start.append(count)
                    temp = int(val.strip())
                    count += temp
                    self.h_res_end.append(count)
            print("    End reading excluded list")
            file.close()
            self.is_initialized = 1
            if self.is_initialized:
                self.read_res_mass()

class trajectory_output:
    """trajectory_output"""

    def __init__(self, controller, md_info):
        self.current_crd_synchronized_step = 0
        self.is_molecule_map_output = 0
        if "molecule_map_output" in controller.Command_Set:
            self.is_molecule_map_output = int(controller.Command_Set["molecule_map_output"])
        self.amber_irest = -1
        self.write_trajectory_interval = 1000 if "write_information_interval" not in controller.Command_Set else int(
            controller.Command_Set["write_information_interval"])
        self.write_restart_file_interval = self.write_trajectory_interval if "write_restart_file_interval" not in \
                                                                             controller.Command_Set else \
            int(controller.Command_Set["write_restart_file_interval"])