#
# Princeton University licenses this file to You 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.
#
#
# ****************************************   DISTRIBUTION FUNCTIONS   **************************************************
"""

* `NormalDist`
* `UniformToNormalDist`
* `ExponentialDist`
* `UniformDist`
* `GammaDist`
* `WaldDist`

Overview
--------

Functions that return one or more samples from a distribution.

"""
from enum import IntEnum

import numpy as np
import typecheck as tc

from psyneulink.core.components.functions.function import Function_Base, FunctionError
from psyneulink.core.globals.keywords import \
    ADDITIVE_PARAM, DIST_FUNCTION_TYPE, BETA, DIST_MEAN, DIST_SHAPE, DRIFT_DIFFUSION_ANALYTICAL_FUNCTION, \
    EXPONENTIAL_DIST_FUNCTION, GAMMA_DIST_FUNCTION, HIGH, LOW, MULTIPLICATIVE_PARAM, NOISE, NORMAL_DIST_FUNCTION, \
    SCALE, STANDARD_DEVIATION, THRESHOLD, UNIFORM_DIST_FUNCTION, WALD_DIST_FUNCTION
from psyneulink.core.globals.context import ContextFlags
from psyneulink.core.globals.utilities import parameter_spec
from psyneulink.core.globals.preferences.basepreferenceset import is_pref_set

from psyneulink.core.globals.parameters import Parameter

__all__ = [
    'DistributionFunction', 'DRIFT_RATE', 'DRIFT_RATE_VARIABILITY', 'DriftDiffusionAnalytical', 'ExponentialDist',
    'GammaDist', 'NON_DECISION_TIME', 'NormalDist', 'STARTING_POINT', 'STARTING_POINT_VARIABILITY',
    'THRESHOLD_VARIABILITY', 'UniformDist', 'UniformToNormalDist', 'WaldDist',
]

class DistributionFunction(Function_Base):
    componentType = DIST_FUNCTION_TYPE

class NormalDist(DistributionFunction):
    """
    NormalDist(                      \
             mean=0.0,               \
             standard_deviation=1.0, \
             params=None,            \
             owner=None,             \
             prefs=None              \
             )

    .. _NormalDist:

    Return a random sample from a normal distribution using numpy.random.normal;

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `standard_deviation <NormalDist.standard_deviation>`
    | *ADDITIVE_PARAM:* `mean <NormalDist.mean>`
    |

    Arguments
    ---------

    mean : float : default 0.0
        The mean or center of the normal distribution

    standard_deviation : float : default 1.0
        Standard deviation of the normal distribution. Must be > 0.0

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    Attributes
    ----------

    mean : float : default 0.0
        The mean or center of the normal distribution.

    standard_deviation : float : default 1.0
        Standard deviation of the normal distribution; if it is 0.0, returns `mean <NormalDist.mean>`.

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    """

    componentName = NORMAL_DIST_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                mean
                    see `mean <NormalDist.mean>`

                    :default value: 0.0
                    :type: float

                standard_deviation
                    see `standard_deviation <NormalDist.standard_deviation>`

                    :default value: 1.0
                    :type: float

        """
        mean = Parameter(0.0, modulable=True, aliases=[ADDITIVE_PARAM])
        standard_deviation = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 mean=0.0,
                 standard_deviation=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(mean=mean,
                                                  standard_deviation=standard_deviation,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _validate_params(self, request_set, target_set=None, context=None):
        super()._validate_params(request_set=request_set, target_set=target_set, context=context)

        if STANDARD_DEVIATION in target_set:
            if target_set[STANDARD_DEVIATION] < 0.0:
                raise FunctionError("The standard_deviation parameter ({}) of {} must be greater than zero.".
                                    format(target_set[STANDARD_DEVIATION], self.name))

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):
        mean = self.get_current_function_param(DIST_MEAN, context)
        standard_deviation = self.get_current_function_param(STANDARD_DEVIATION, context)

        result = np.random.normal(mean, standard_deviation)

        return self.convert_output_type(result)

class UniformToNormalDist(DistributionFunction):
    """
    UniformToNormalDist(             \
             mean=0.0,               \
             standard_deviation=1.0, \
             params=None,            \
             owner=None,             \
             prefs=None              \
             )

    .. _UniformToNormalDist:

    Return a random sample from a normal distribution using first np.random.rand(1) to generate a sample from a uniform
    distribution, and then converting that sample to a sample from a normal distribution with the following equation:

    .. math::

        normal\\_sample = \\sqrt{2} \\cdot standard\\_dev \\cdot scipy.special.erfinv(2 \\cdot uniform\\_sample - 1)  + mean

    The uniform --> normal conversion allows for a more direct comparison with MATLAB scripts.

    .. note::

        This function requires `SciPy <https://pypi.python.org/pypi/scipy>`_.

    (https://github.com/jonasrauber/randn-matlab-python)

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `standard_deviation <UniformToNormalDist.standard_deviation>`
    | *ADDITIVE_PARAM:* `mean <UniformToNormalDist.mean>`
    |

    Arguments
    ---------

    mean : float : default 0.0
        The mean or center of the normal distribution

    standard_deviation : float : default 1.0
        Standard deviation of the normal distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    Attributes
    ----------

    mean : float : default 0.0
        The mean or center of the normal distribution

    standard_deviation : float : default 1.0
        Standard deviation of the normal distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    """

    componentName = NORMAL_DIST_FUNCTION

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                variable
                    see `variable <UniformToNormalDist.variable>`

                    :default value: numpy.array([0])
                    :type: numpy.ndarray
                    :read only: True

                mean
                    see `mean <UniformToNormalDist.mean>`

                    :default value: 0.0
                    :type: float

                standard_deviation
                    see `standard_deviation <UniformToNormalDist.standard_deviation>`

                    :default value: 1.0
                    :type: float

        """
        variable = Parameter(np.array([0]), read_only=True, pnl_internal=True, constructor_argument='default_variable')
        mean = Parameter(0.0, modulable=True, aliases=[ADDITIVE_PARAM])
        standard_deviation = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 mean=0.0,
                 standard_deviation=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(mean=mean,
                                                  standard_deviation=standard_deviation,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):

        try:
            from scipy.special import erfinv
        except:
            raise FunctionError("The UniformToNormalDist function requires the SciPy package.")

        mean = self.get_current_function_param(DIST_MEAN, context)
        standard_deviation = self.get_current_function_param(STANDARD_DEVIATION, context)

        sample = np.random.rand(1)[0]
        result = ((np.sqrt(2) * erfinv(2 * sample - 1)) * standard_deviation) + mean

        return self.convert_output_type(result)

class ExponentialDist(DistributionFunction):
    """
    ExponentialDist(                \
             beta=1.0,              \
             params=None,           \
             owner=None,            \
             prefs=None             \
             )

    .. _ExponentialDist:

    Return a random sample from a exponential distribution using numpy.random.exponential

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `beta <ExponentialDist.beta>`
    |

    Arguments
    ---------

    beta : float : default 1.0
        The scale parameter of the exponential distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    Attributes
    ----------

    beta : float : default 1.0
        The scale parameter of the exponential distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    """
    componentName = EXPONENTIAL_DIST_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                beta
                    see `beta <ExponentialDist.beta>`

                    :default value: 1.0
                    :type: float

        """
        beta = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 beta=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(beta=beta,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):

        beta = self.get_current_function_param(BETA, context)
        result = np.random.exponential(beta)

        return self.convert_output_type(result)

class UniformDist(DistributionFunction):
    """
    UniformDist(                      \
             low=0.0,             \
             high=1.0,             \
             params=None,           \
             owner=None,            \
             prefs=None             \
             )

    .. _UniformDist:

    Return a random sample from a uniform distribution using numpy.random.uniform

    Arguments
    ---------

    low : float : default 0.0
        Lower bound of the uniform distribution

    high : float : default 1.0
        Upper bound of the uniform distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    Attributes
    ----------

    low : float : default 0.0
        Lower bound of the uniform distribution

    high : float : default 1.0
        Upper bound of the uniform distribution

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    """
    componentName = UNIFORM_DIST_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                high
                    see `high <UniformDist.high>`

                    :default value: 1.0
                    :type: float

                low
                    see `low <UniformDist.low>`

                    :default value: 0.0
                    :type: float

        """
        low = Parameter(0.0, modulable=True)
        high = Parameter(1.0, modulable=True)

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 low=0.0,
                 high=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(low=low,
                                                  high=high,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):

        low = self.get_current_function_param(LOW, context)
        high = self.get_current_function_param(HIGH, context)
        result = np.random.uniform(low, high)

        return self.convert_output_type(result)

class GammaDist(DistributionFunction):
    """
    GammaDist(\
             scale=1.0,\
             dist_shape=1.0,\
             params=None,\
             owner=None,\
             prefs=None\
             )

    .. _GammaDist:

    Return a random sample from a gamma distribution using numpy.random.gamma

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `scale <GammaDist.scale>`
    | *ADDITIVE_PARAM:* `dist_shape <GammaDist.dist_shape>`
    |

    Arguments
    ---------

    scale : float : default 1.0
        The scale of the gamma distribution. Should be greater than zero.

    dist_shape : float : default 1.0
        The shape of the gamma distribution. Should be greater than zero.

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    Attributes
    ----------

    scale : float : default 1.0
        The scale of the gamma distribution. Should be greater than zero.

    dist_shape : float : default 1.0
        The shape of the gamma distribution. Should be greater than zero.

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    """

    componentName = GAMMA_DIST_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                dist_shape
                    see `dist_shape <GammaDist.dist_shape>`

                    :default value: 1.0
                    :type: float

                scale
                    see `scale <GammaDist.scale>`

                    :default value: 1.0
                    :type: float

        """
        scale = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])
        dist_shape = Parameter(1.0, modulable=True, aliases=[ADDITIVE_PARAM])

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 scale=1.0,
                 dist_shape=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(scale=scale,
                                                  dist_shape=dist_shape,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):

        scale = self.get_current_function_param(SCALE, context)
        dist_shape = self.get_current_function_param(DIST_SHAPE, context)

        result = np.random.gamma(dist_shape, scale)

        return self.convert_output_type(result)

class WaldDist(DistributionFunction):
    """
     WaldDist(             \
              scale=1.0,\
              mean=1.0,\
              params=None,\
              owner=None,\
              prefs=None\
              )

     .. _WaldDist:

     Return a random sample from a Wald distribution using numpy.random.wald

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `scale <WaldDist.scale>`
    | *ADDITIVE_PARAM:* `mean <WaldDist.mean>`
    |

     Arguments
     ---------

     scale : float : default 1.0
         Scale parameter of the Wald distribution. Should be greater than zero.

     mean : float : default 1.0
         Mean of the Wald distribution. Should be greater than or equal to zero.

     params : Dict[param keyword: param value] : default None
         a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
         function.  Values specified for parameters in the dictionary override any assigned to those parameters in
         arguments of the constructor.

     owner : Component
         `component <Component>` to which to assign the Function.

     prefs : PreferenceSet or specification dict : default Function.classPreferences
         the `PreferenceSet` for the Function. If it is not specified, a default is assigned using `classPreferences`
         defined in __init__.py (see :doc:`PreferenceSet <LINK>` for details).

     Attributes
     ----------

     scale : float : default 1.0
         Scale parameter of the Wald distribution. Should be greater than zero.

     mean : float : default 1.0
         Mean of the Wald distribution. Should be greater than or equal to zero.

     params : Dict[param keyword: param value] : default None
         a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
         function.  Values specified for parameters in the dictionary override any assigned to those parameters in
         arguments of the constructor.

     owner : Component
         `component <Component>` to which to assign the Function.

     prefs : PreferenceSet or specification dict : default Function.classPreferences
         the `PreferenceSet` for the Function. If it is not specified, a default is assigned using `classPreferences`
         defined in __init__.py (see :doc:`PreferenceSet <LINK>` for details).

     """

    componentName = WALD_DIST_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                mean
                    see `mean <WaldDist.mean>`

                    :default value: 1.0
                    :type: float

                scale
                    see `scale <WaldDist.scale>`

                    :default value: 1.0
                    :type: float

        """
        scale = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])
        mean = Parameter(1.0, modulable=True, aliases=[ADDITIVE_PARAM])

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 scale=1.0,
                 mean=1.0,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None):
        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(scale=scale,
                                                  mean=mean,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):

        scale = self.get_current_function_param(SCALE, context)
        mean = self.get_current_function_param(DIST_MEAN, context)

        result = np.random.wald(mean, scale)

        return self.convert_output_type(result)

# Note:  For any of these that correspond to args, value must match the name of the corresponding arg in __init__()
DRIFT_RATE = 'drift_rate'
DRIFT_RATE_VARIABILITY = 'DDM_DriftRateVariability'
THRESHOLD_VARIABILITY = 'DDM_ThresholdRateVariability'
STARTING_POINT = 'starting_point'
STARTING_POINT_VARIABILITY = "DDM_StartingPointVariability"
NON_DECISION_TIME = 't0'

def _DriftDiffusionAnalytical_bias_getter(owning_component=None, context=None):
    starting_point = owning_component.parameters.starting_point._get(context)
    threshold = owning_component.parameters.threshold._get(context)
    try:
        return (starting_point + threshold) / (2 * threshold)
    except TypeError:
        return None

# QUESTION: IF VARIABLE IS AN ARRAY, DOES IT RETURN AN ARRAY FOR EACH RETURN VALUE (RT, ER, ETC.)
class DriftDiffusionAnalytical(DistributionFunction):  # -------------------------------------------------------------------------------
    """
    DriftDiffusionAnalytical(   \
        default_variable=None,  \
        drift_rate=1.0,         \
        threshold=1.0,          \
        starting_point=0.0,     \
        t0=0.2                  \
        noise=0.5,              \
        params=None,            \
        owner=None,             \
        prefs=None              \
        )

    .. _DriftDiffusionAnalytical:

    Return terminal value of decision variable, mean accuracy, and mean response time computed analytically for the
    drift diffusion process as described in `Bogacz et al (2006) <https://www.ncbi.nlm.nih.gov/pubmed/17014301>`_.

    *Modulatory Parameters:*

    | *MULTIPLICATIVE_PARAM:* `drift_rate <DriftDiffusionAnalytical.drift_rate>`
    | *ADDITIVE_PARAM:* `starting_point <DriftDiffusionAnalytical.starting_point>`
    |

    Arguments
    ---------

    default_variable : number, list or array : default class_defaults.variable
        specifies a template for decision variable(s);  if it is list or array, a separate solution is computed
        independently for each element.

    drift_rate : float, list or 1d array : default 1.0
        specifies the drift_rate of the drift diffusion process.  If it is a list or array,
        it must be the same length as `default_variable <DriftDiffusionAnalytical.default_variable>`.

    threshold : float, list or 1d array : default 1.0
        specifies the threshold (boundary) of the drift diffusion process.  If it is a list or array,
        it must be the same length as `default_variable <DriftDiffusionAnalytical.default_variable>`.

    starting_point : float, list or 1d array : default 1.0
        specifies the initial value of the decision variable for the drift diffusion process.  If it is a list or
        array, it must be the same length as `default_variable <DriftDiffusionAnalytical.default_variable>`.

    noise : float, list or 1d array : default 0.0
        specifies the noise term (corresponding to the diffusion component) of the drift diffusion process.
        If it is a float, it must be a number from 0 to 1.  If it is a list or array, it must be the same length as
        `default_variable <DriftDiffusionAnalytical.default_variable>` and all elements must be floats from 0 to 1.

    t0 : float, list or 1d array : default 0.2
        specifies the non-decision time for solution. If it is a float, it must be a number from 0 to 1.  If it is a
        list or array, it must be the same length as  `default_variable <DriftDiffusionAnalytical.default_variable>` and all
        elements must be floats from 0 to 1.

    params : Dict[param keyword: param value] : default None
        a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
        function.  Values specified for parameters in the dictionary override any assigned to those parameters in
        arguments of the constructor.

    owner : Component
        `component <Component>` to which to assign the Function.

    name : str : default see `name <Function.name>`
        specifies the name of the Function.

    prefs : PreferenceSet or specification dict : default Function.classPreferences
        specifies the `PreferenceSet` for the Function (see `prefs <Function_Base.prefs>` for details).

    shenhav_et_al_compat_mode: bool : default False
        whether Shenhav et al. compatibility mode is set. See shenhav_et_al_compat_mode property.

    Attributes
    ----------

    variable : number or 1d array
        holds initial value assigned to :keyword:`default_variable` argument;
        ignored by `function <BogaczEtal.function>`.

    drift_rate : float or 1d array
        determines the drift component of the drift diffusion process.

    threshold : float or 1d array
        determines the threshold (boundary) of the drift diffusion process (i.e., at which the integration
        process is assumed to terminate).

    starting_point : float or 1d array
        determines the initial value of the decision variable for the drift diffusion process.

    noise : float or 1d array
        determines the diffusion component of the drift diffusion process (used to specify the variance of a
        Gaussian random process).

    t0 : float or 1d array
        determines the assumed non-decision time to determine the response time returned by the solution.

    bias : float or 1d array
        normalized starting point:
        (`starting_point <DriftDiffusionAnalytical.starting_point>` + `threshold <DriftDiffusionAnalytical.threshold>`) /
        (2 * `threshold <DriftDiffusionAnalytical.threshold>`)

    owner : Component
        `component <Component>` to which the Function has been assigned.

    name : str
        the name of the Function; if it is not specified in the **name** argument of the constructor, a
        default is assigned by FunctionRegistry (see `Naming` for conventions used for default and duplicate names).

    prefs : PreferenceSet or specification dict : Function.classPreferences
        the `PreferenceSet` for function; if it is not specified in the **prefs** argument of the Function's
        constructor, a default is assigned using `classPreferences` defined in __init__.py (see :doc:`PreferenceSet
        <LINK>` for details).

    """

    componentName = DRIFT_DIFFUSION_ANALYTICAL_FUNCTION

    paramClassDefaults = Function_Base.paramClassDefaults.copy()

    class Parameters(DistributionFunction.Parameters):
        """
            Attributes
            ----------

                bias
                    see `bias <DriftDiffusionAnalytical.bias>`

                    :default value: 0.5
                    :type: float
                    :read only: True

                drift_rate
                    see `drift_rate <DriftDiffusionAnalytical.drift_rate>`

                    :default value: 1.0
                    :type: float

                noise
                    see `noise <DriftDiffusionAnalytical.noise>`

                    :default value: 0.5
                    :type: float

                starting_point
                    see `starting_point <DriftDiffusionAnalytical.starting_point>`

                    :default value: 0.0
                    :type: float

                t0
                    see `t0 <DriftDiffusionAnalytical.t0>`

                    :default value: 0.2
                    :type: float

                threshold
                    see `threshold <DriftDiffusionAnalytical.threshold>`

                    :default value: 1.0
                    :type: float

        """
        drift_rate = Parameter(1.0, modulable=True, aliases=[MULTIPLICATIVE_PARAM])
        starting_point = Parameter(0.0, modulable=True, aliases=[ADDITIVE_PARAM])
        threshold = Parameter(1.0, modulable=True)
        noise = Parameter(0.5, modulable=True)
        t0 = .200
        bias = Parameter(0.5, read_only=True, getter=_DriftDiffusionAnalytical_bias_getter)

    @tc.typecheck
    def __init__(self,
                 default_variable=None,
                 drift_rate: parameter_spec = 1.0,
                 starting_point: parameter_spec = 0.0,
                 threshold: parameter_spec = 1.0,
                 noise: parameter_spec = 0.5,
                 t0: parameter_spec = .200,
                 params=None,
                 owner=None,
                 prefs: is_pref_set = None,
                 shenhav_et_al_compat_mode=False):

        self._shenhav_et_al_compat_mode = shenhav_et_al_compat_mode

        # Assign args to params and functionParams dicts
        params = self._assign_args_to_param_dicts(drift_rate=drift_rate,
                                                  starting_point=starting_point,
                                                  threshold=threshold,
                                                  noise=noise,
                                                  t0=t0,
                                                  params=params)

        super().__init__(default_variable=default_variable,
                         params=params,
                         owner=owner,
                         prefs=prefs,
                         )

    @property
    def output_type(self):
        return self._output_type

    @output_type.setter
    def output_type(self, value):
        # disabled because it happens during normal execution, may be confusing
        # warnings.warn('output_type conversion disabled for {0}'.format(self.__class__.__name__))
        self._output_type = None

    @property
    def shenhav_et_al_compat_mode(self):
        """
        Get whether the function is set to Shenhav et al. compatibility mode. This mode allows
        the analytic computations of mean error rate and reaction time to match exactly the
        computations made in the MATLAB DDM code (Matlab/ddmSimFRG.m). These compatibility changes
        should only effect edges cases that involve the following cases:

            - Floating point overflows and underflows are ignored when computing mean RT and mean ER
            - Exponential expressions used in cacluating mean RT and mean ER are bounded by 1e-12 to 1e12.
            - Decision time is not permitted to be negative and will be set to 0 in these cases. Thus RT
              will be RT = non-decision-time in these cases.

        Returns
        -------
        Shenhav et al. compatible mode setting : (bool)

        """
        return self._shenhav_et_al_compat_mode

    @shenhav_et_al_compat_mode.setter
    def shenhav_et_al_compat_mode(self, value):
        """
        Set whether the function is set to Shenhav et al. compatibility mode. This mode allows
        the analytic computations of mean error rate and reaction time to match exactly the
        computations made in the MATLAB DDM code (Matlab/ddmSimFRG.m). These compatibility chages
        should only effect edges cases that involve the following cases:

            - Floating point overflows and underflows are ignored when computing mean RT and mean ER
            - Exponential expressions used in cacluating mean RT and mean ER are bounded by 1e-12 to 1e12.
            - Decision time is not permitted to be negative and will be set to 0 in these cases. Thus RT
              will be RT = non-decision-time in these cases.

        Arguments
        ---------

        value : bool
            Set True to turn on Shenhav et al. compatibility mode, False for off.
        """
        self._shenhav_et_al_compat_mode = value

    def _function(self,
                 variable=None,
                 context=None,
                 params=None,
                 ):
        """
        Return: terminal value of decision variable (equal to threshold), mean accuracy (error rate; ER) and mean
        response time (RT)

        Arguments
        ---------

        variable : 2d array
            ignored.

        params : Dict[param keyword: param value] : default None
            a `parameter dictionary <ParameterPort_Specification>` that specifies the parameters for the
            function.  Values specified for parameters in the dictionary override any assigned to those parameters in
            arguments of the constructor.

        Returns
        -------
        Decision variable, mean ER, mean RT : (float, float, float)

        """

        attentional_drift_rate = float(self.get_current_function_param(DRIFT_RATE, context))
        stimulus_drift_rate = float(variable)
        drift_rate = attentional_drift_rate * stimulus_drift_rate
        threshold = self.get_current_function_param(THRESHOLD, context)
        starting_point = float(self.get_current_function_param(STARTING_POINT, context))
        noise = float(self.get_current_function_param(NOISE, context))
        t0 = float(self.get_current_function_param(NON_DECISION_TIME, context))

        # drift_rate = float(self.drift_rate) * float(variable)
        # threshold = float(self.threshold)
        # starting_point = float(self.starting_point)
        # noise = float(self.noise)
        # t0 = float(self.t0)

        bias = (starting_point + threshold) / (2 * threshold)

        # Prevents div by 0 issue below:
        if bias <= 0:
            bias = 1e-8
        if bias >= 1:
            bias = 1 - 1e-8

        # drift_rate close to or at 0 (avoid float comparison)
        if np.abs(drift_rate) < 1e-8:
            # back to absolute bias in order to apply limit
            bias_abs = bias * 2 * threshold - threshold
            # use expression for limit a->0 from Srivastava et al. 2016
            rt = t0 + (threshold ** 2 - bias_abs ** 2) / (noise ** 2)
            er = (threshold - bias_abs) / (2 * threshold)
        else:
            drift_rate_normed = np.abs(drift_rate)
            ztilde = threshold / drift_rate_normed
            atilde = (drift_rate_normed / noise) ** 2

            is_neg_drift = drift_rate < 0
            bias_adj = (is_neg_drift == 1) * (1 - bias) + (is_neg_drift == 0) * bias
            y0tilde = ((noise ** 2) / 2) * np.log(bias_adj / (1 - bias_adj))
            if np.abs(y0tilde) > threshold:
                # First difference between Shenhav et al. DDM code and PNL's.
                if self.shenhav_et_al_compat_mode:
                    y0tilde = -1 * (y0tilde < 0) * threshold + (y0tilde >=0 ) * threshold
                else:
                    y0tilde = -1 * (is_neg_drift == 1) * threshold + (is_neg_drift == 0) * threshold

            x0tilde = y0tilde / drift_rate_normed

            # Whether we should ignore or raise floating point over and underflow exceptions.
            # Shenhav et al. MATLAB code ignores them.
            ignore_or_raise = "raise"
            if self.shenhav_et_al_compat_mode:
                ignore_or_raise = "ignore"

            with np.errstate(over=ignore_or_raise, under=ignore_or_raise):
                try:
                    # Lets precompute these common sub-expressions
                    exp_neg2_x0tilde_atilde = np.exp(-2 * x0tilde * atilde)
                    exp_2_ztilde_atilde = np.exp(2 * ztilde * atilde)
                    exp_neg2_ztilde_atilde = np.exp(-2 * ztilde * atilde)

                    if self.shenhav_et_al_compat_mode:
                        exp_neg2_x0tilde_atilde = np.nanmax([1e-12, exp_neg2_x0tilde_atilde])
                        exp_2_ztilde_atilde = np.nanmin([1e12, exp_2_ztilde_atilde])
                        exp_neg2_ztilde_atilde = np.nanmax([1e-12, exp_neg2_ztilde_atilde])

                    rt = ztilde * np.tanh(ztilde * atilde) + \
                         ((2 * ztilde * (1 - exp_neg2_x0tilde_atilde)) / (
                                 exp_2_ztilde_atilde - exp_neg2_ztilde_atilde) - x0tilde)
                    er = 1 / (1 + exp_2_ztilde_atilde) - \
                         ((1 - exp_neg2_x0tilde_atilde) / (exp_2_ztilde_atilde - exp_neg2_ztilde_atilde))

                    # Fail safe to prevent negative mean RT's. Shenhav et al. do this.
                    if self.shenhav_et_al_compat_mode:
                        if rt < 0:
                            rt = 0

                    rt = rt + t0

                except FloatingPointError:
                    # Per Mike Shvartsman:
                    # If ±2*ztilde*atilde (~ 2*z*a/(c^2) gets very large, the diffusion vanishes relative to drift
                    # and the problem is near-deterministic. Without diffusion, error rate goes to 0 or 1
                    # depending on the sign of the drift, and so decision time goes to a point mass on z/a – x0, and
                    # generates a "RuntimeWarning: overflow encountered in exp"
                    er = 0
                    rt = ztilde / atilde - x0tilde + t0

            # This last line makes it report back in terms of a fixed reference point
            #    (i.e., closer to 1 always means higher p(upper boundary))
            # If you comment this out it will report errors in the reference frame of the drift rate
            #    (i.e., reports p(upper) if drift is positive, and p(lower if drift is negative)
            er = (is_neg_drift == 1) * (1 - er) + (is_neg_drift == 0) * (er)

        # Compute moments (mean, variance, skew) of condiational response time distributions
        moments = DriftDiffusionAnalytical._compute_conditional_rt_moments(drift_rate, noise, threshold, bias, t0)

        return rt, er, \
               moments['mean_rt_plus'], moments['var_rt_plus'], moments['skew_rt_plus'], \
               moments['mean_rt_minus'], moments['var_rt_minus'], moments['skew_rt_minus']

    @staticmethod
    def _compute_conditional_rt_moments(drift_rate, noise, threshold, starting_point, t0):
        """
        This is a helper function for computing the conditional decison time moments for the DDM.
        It is based completely off of Matlab\\DDMFunctions\\ddm_metrics_cond_Mat.m.

        :param drift_rate: The drift rate of the DDM
        :param noise: The diffusion rate.
        :param threshold: The symmetric threshold of the DDM
        :param starting_point: The initial condition.
        :param t0: The non decision time.
        :return: A dictionary containing the following key value pairs:
         mean_rt_plus: The mean RT of positive responses.
         mean_rt_minus: The mean RT of negative responses.
         var_rt_plus: The variance of RT of positive responses.
         var_rt_minus: The variance of RT of negative responses.
         skew_rt_plus: The skew of RT of positive responses.
         skew_rt_minus: The skew of RT of negative responses.
        """

        #  transform starting point to be centered at 0
        starting_point = (starting_point - 0.5) * 2.0 * threshold

        if abs(drift_rate) < 0.01:
            drift_rate = 0.01

        X = drift_rate * starting_point / noise**2
        Z = drift_rate * threshold / noise**2

        X = max(-100, min(100, X))

        Z = max(-100, min(100, Z))

        if abs(Z) < 0.0001:
            Z = 0.0001

        def coth(x):
            return 1 / np.tanh(x)

        def csch(x):
            return 1 / np.sinh(x)

        moments = {}

        # Lets ignore any divide by zeros we get or NaN errors. This will allow the NaN's to propogate.
        with np.errstate(divide='ignore', invalid='ignore'):
            moments["mean_rt_plus"] = noise**2. / (drift_rate**2) * (2 * Z * coth(2 * Z) - (X + Z) * coth(X + Z))

            moments["mean_rt_minus"] = noise**2. / (drift_rate**2) * (2 * Z * coth(2 * Z) - (-X + Z) * coth(-X + Z))

            moments["var_rt_plus"] = noise**4. / (drift_rate**4) * \
                              (4 * Z**2. * (csch(2 * Z))**2 + 2 * Z * coth(2 * Z) - (Z + X)**2. *
                               (csch(Z + X))**2 - (Z + X) * coth(Z + X))

            moments["var_rt_minus"] = noise**4. / (drift_rate**4) * \
                               (4 * Z**2. * (csch(2 * Z)) ** 2 + 2 * Z * coth(2 * Z) - (Z - X)**2. *
                                (csch(Z - X))**2 - (Z - X) * coth(Z - X))

            moments["skew_rt_plus"] = noise**6. / (drift_rate** 6) * \
                               (12 * Z**2. * (csch(2 * Z))**2 + 16 * Z**3. * coth(2 * Z) * (csch(2 * Z))**2 +
                                6 * Z * coth(2 * Z) - 3 * (Z + X)**2. * (csch(Z + X))**2 -
                                2 * (Z + X)**3. * coth(Z + X) * (csch(Z + X))**2 - 3 * (Z + X) * coth(Z + X))

            moments["skew_rt_minus"] = noise**6. / (drift_rate**6) * \
                                (12 * Z**2. * (csch(2 * Z))**2 + 16 * Z**3. * coth(2 * Z) *
                                 (csch(2 * Z))**2 + 6 * Z * coth(2 * Z) - 3 * (Z - X)**2. *
                                 (csch(Z - X))**2 - 2 * (Z - X)**3. * coth(Z - X) *
                                 (csch(Z - X))**2 - 3 * (Z - X) * coth(Z - X))

            # divide third central moment by var_rt**1.5 to get skewness
            moments['skew_rt_plus'] /= moments['var_rt_plus']**1.5
            moments['skew_rt_minus'] /= moments['var_rt_minus']**1.5

            # Add the non-decision time to the mean RTs
            moments['mean_rt_plus'] += t0
            moments['mean_rt_minus'] += t0

        return moments

    def derivative(self, output=None, input=None, context=None):
        """
        derivative(output, input)

        Calculate the derivative of :math:`\\frac{1}{reward rate}` with respect to the threshold (**output** arg)
        and drift_rate (**input** arg).  Reward rate (:math:`RR`) is assumed to be:

            :math:`RR = delay_{ITI} + \\frac{Z}{A} + ED`;

        the derivative of :math:`\\frac{1}{RR}` with respect to the `threshold <DriftDiffusionAnalytical.threshold>` is:

            :math:`\\frac{1}{A} - \\frac{E}{A} - 2\\frac{A}{c^2}ED`;

        and the derivative of 1/RR with respect to the `drift_rate <DriftDiffusionAnalytical.drift_rate>` is:

            :math:`-\\frac{Z}{A^2} + \\frac{Z}{A^2}E - \\frac{2Z}{c^2}ED`

        where:

            *A* = `drift_rate <DriftDiffusionAnalytical.drift_rate>`,

            *Z* = `threshold <DriftDiffusionAnalytical.threshold>`,

            *c* = `noise <DriftDiffusionAnalytical.noise>`,

            *E* = :math:`e^{-2\\frac{ZA}{c^2}}`,

            *D* = :math:`delay_{ITI} + delay_{penalty} - \\frac{Z}{A}`,

            :math:`delay_{ITI}` is the intertrial interval and :math:`delay_{penalty}` is a penalty delay.

        Returns
        -------

        derivatives :  List[float, float)
            of :math:`\\frac{1}{RR}` with respect to `threshold <DriftDiffusionAnalytical.threshold>` and `drift_rate
            <DriftDiffusionAnalytical.drift_rate>`.

        """
        Z = output or self.get_current_function_param(THRESHOLD, context)
        A = input or self.get_current_function_param(DRIFT_RATE, context)
        c = self.get_current_function_param(NOISE, context)
        c_sq = c ** 2
        E = np.exp(-2 * Z * A / c_sq)
        D_iti = 0
        D_pen = 0
        D = D_iti + D_pen
        # RR =  1/(D_iti + Z/A + (E*D))

        dRR_dZ = 1 / A + E / A + (2 * A / c_sq) * E * D
        dRR_dA = -Z / A ** 2 + (Z / A ** 2) * E - (2 * Z / c_sq) * E * D

        return [dRR_dZ, dRR_dA]


import logging

logger = logging.getLogger()

logging.basicConfig(
    filename='application-log.log',
    filemode='w',
    format='%(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

CITIES = {
    "MOSCOW": "https://code.s3.yandex.net/async-module/moscow-response.json",
    "PARIS": "https://code.s3.yandex.net/async-module/paris-response.json",
    "LONDON": "https://code.s3.yandex.net/async-module/london-response.json",
    "BERLIN": "https://code.s3.yandex.net/async-module/berlin-response.json",
    "BEIJING": "https://code.s3.yandex.net/async-module/beijing-response.json",
    "KAZAN": "https://code.s3.yandex.net/async-module/kazan-response.json",
    "SPETERSBURG": "https://code.s3.yandex.net/async-module/spetersburg-response.json",
    "VOLGOGRAD": "https://code.s3.yandex.net/async-module/volgograd-response.json",
    "NOVOSIBIRSK": "https://code.s3.yandex.net/async-module/novosibirsk-response.json",
    "KALININGRAD": "https://code.s3.yandex.net/async-module/kaliningrad-response.json",
    "ABUDHABI": "https://code.s3.yandex.net/async-module/abudhabi-response.json",
    "WARSZAWA": "https://code.s3.yandex.net/async-module/warszawa-response.json",
    "BUCHAREST": "https://code.s3.yandex.net/async-module/bucharest-response.json",
    "ROMA": "https://code.s3.yandex.net/async-module/roma-response.json",
    "CAIRO": "https://code.s3.yandex.net/async-module/cairo-response.json",
}
ERR_MESSAGE_TEMPLATE = "Something wrong. Please contact with mentor."

MIN_MAJOR_PYTHON_VER = 3
MIN_MINOR_PYTHON_VER = 9

MIN_TIME_IN_EACH_DATE = 10
MAX_TIME_IN_EACH_DATE = 19

GOOD_WEATHER = [
    'clear', 'partly-cloudy', 'cloudy', 'overcast', 'drizzle'
]

DF_TEMPLATE = [
    {
        'Погода': 'Температура, среднее'
    },
    {
        'Погода': 'Без осадков, часов'
    }
]

MSG_NOT_ENOUGH_DATA = 'Недостаточно данных по часам для рассчета в городе "{city}". Дата: "{date}" исключена.'
MSG_RECOMMENDED_CITIES = 'Вам рекомендовано отправиться в такие города, как "{cities}"'
MSG_QUEUT_PUT = 'В очередь добавлен "{item}"'
MSG_CLEAR_QUEUE = 'Очередь пуста'
MSG_GET_ITEM_QUEUE = 'Запись "{item}" получена из очереди'
MSG_CSV_SAVED = 'Файл csv успешно сохранен'
MSG_FINAL = 'Метод analyze отработал, результат: "{cities}"'

def check_python_version():
    import sys

    if (
            sys.version_info.major < MIN_MAJOR_PYTHON_VER
            or sys.version_info.minor < MIN_MINOR_PYTHON_VER
    ):
        raise Exception(
            'Please use python version >= {}.{}'.format(
                MIN_MAJOR_PYTHON_VER, MIN_MINOR_PYTHON_VER
            )
        )


import logging
from typing import Any, Dict, Iterable, Optional, Tuple, Union

import torch
import torch_optimizer as optim
from einops import rearrange, repeat
from nowcasting_dataset.consts import (
    DATETIME_FEATURE_NAMES,
    NWP_DATA,
    NWP_X_COORDS,
    NWP_Y_COORDS,
    SATELLITE_DATA,
    SATELLITE_DATETIME_INDEX,
    SATELLITE_X_COORDS,
    SATELLITE_Y_COORDS,
    TOPOGRAPHIC_DATA,
    TOPOGRAPHIC_X_COORDS,
    TOPOGRAPHIC_Y_COORDS,
)
from nowcasting_utils.models.base import BaseModel, register_model
from nowcasting_utils.models.loss import get_loss
from perceiver_pytorch import MultiPerceiver
from perceiver_pytorch.decoders import ImageDecoder
from perceiver_pytorch.encoders import ImageEncoder
from perceiver_pytorch.modalities import InputModality
from perceiver_pytorch.queries import LearnableQuery
from perceiver_pytorch.utils import encode_position
from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR

logger = logging.getLogger("satflow.model")
logger.setLevel(logging.WARN)

@register_model
class Perceiver(BaseModel):
    def __init__(
        self,
        input_channels: int = 22,
        sat_channels: int = 12,
        nwp_channels: int = 10,
        base_channels: int = 1,
        forecast_steps: int = 48,
        history_steps: int = 6,
        input_size: int = 64,
        lr: float = 5e-4,
        visualize: bool = True,
        max_frequency: float = 4.0,
        depth: int = 6,
        num_latents: int = 256,
        cross_heads: int = 1,
        latent_heads: int = 8,
        cross_dim_heads: int = 8,
        latent_dim: int = 512,
        weight_tie_layers: bool = False,
        decoder_ff: bool = True,
        dim: int = 32,
        logits_dim: int = 100,
        queries_dim: int = 32,
        latent_dim_heads: int = 64,
        loss="mse",
        sin_only: bool = False,
        encode_fourier: bool = True,
        preprocessor_type: Optional[str] = None,
        postprocessor_type: Optional[str] = None,
        encoder_kwargs: Optional[Dict[str, Any]] = None,
        decoder_kwargs: Optional[Dict[str, Any]] = None,
        pretrained: bool = False,
        predict_timesteps_together: bool = False,
        nwp_modality: bool = False,
        datetime_modality: bool = False,
        use_learnable_query: bool = True,
        generate_fourier_features: bool = True,
        temporally_consistent_fourier_features: bool = False,
    ):
        super(BaseModel, self).__init__()
        self.forecast_steps = forecast_steps
        self.input_channels = input_channels
        self.lr = lr
        self.pretrained = pretrained
        self.visualize = visualize
        self.sat_channels = sat_channels
        self.nwp_channels = nwp_channels
        self.output_channels = sat_channels
        self.criterion = get_loss(loss)
        self.input_size = input_size
        self.predict_timesteps_together = predict_timesteps_together
        self.use_learnable_query = use_learnable_query
        self.max_frequency = max_frequency
        self.temporally_consistent_fourier_features = temporally_consistent_fourier_features

        if use_learnable_query:
            self.query = LearnableQuery(
                channel_dim=queries_dim,
                query_shape=(self.forecast_steps, self.input_size, self.input_size)
                if predict_timesteps_together
                else (self.input_size, self.input_size),
                conv_layer="3d",
                max_frequency=max_frequency,
                num_frequency_bands=input_size,
                sine_only=sin_only,
                generate_fourier_features=generate_fourier_features,
            )
        else:
            self.query = None

        # Warn if using frequency is smaller than Nyquist Frequency
        if max_frequency < input_size / 2:
            print(
                f"Max frequency is less than Nyquist frequency, currently set to {max_frequency} while "
                f"the Nyquist frequency for input of size {input_size} is {input_size / 2}"
            )

        # Preprocessor, if desired, on top of the other processing done
        if preprocessor_type is not None:
            if preprocessor_type not in ("conv", "patches", "pixels", "conv1x1", "metnet"):
                raise ValueError("Invalid prep_type!")
            if preprocessor_type == "metnet":
                # MetNet processing
                self.preprocessor = ImageEncoder(
                    crop_size=input_size,
                    prep_type="metnet",
                )
                video_input_channels = (
                    8 * sat_channels
                )  # This is only done on the sat channel inputs
                nwp_input_channels = 8 * nwp_channels
                # If doing it on the base map, then need
                image_input_channels = 4 * base_channels
            else:
                self.preprocessor = ImageEncoder(
                    input_channels=sat_channels,
                    prep_type=preprocessor_type,
                    **encoder_kwargs,
                )
                nwp_input_channels = self.preprocessor.output_channels
                video_input_channels = self.preprocessor.output_channels
                image_input_channels = self.preprocessor.output_channels
        else:
            self.preprocessor = None
            nwp_input_channels = nwp_channels
            video_input_channels = sat_channels
            image_input_channels = base_channels

        # The preprocessor will change the number of channels in the input
        modalities = []
        # Timeseries input
        sat_modality = InputModality(
            name=SATELLITE_DATA,
            input_channels=video_input_channels,
            input_axis=3,  # number of axes, 3 for video
            num_freq_bands=input_size,  # number of freq bands, with original value (2 * K + 1)
            max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is, should be Nyquist frequency (i.e. 112 for 224 input image)
            sin_only=sin_only,  # Whether if sine only for Fourier encoding, TODO test more
            fourier_encode=encode_fourier,  # Whether to encode position with Fourier features
        )
        modalities.append(sat_modality)
        if nwp_modality:
            nwp_modality = InputModality(
                name=NWP_DATA,
                input_channels=nwp_input_channels,
                input_axis=3,  # number of axes, 3 for video
                num_freq_bands=input_size,  # number of freq bands, with original value (2 * K + 1)
                max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is, should be Nyquist frequency (i.e. 112 for 224 input image)
                sin_only=sin_only,  # Whether if sine only for Fourier encoding, TODO test more
                fourier_encode=encode_fourier,  # Whether to encode position with Fourier features
            )
            modalities.append(nwp_modality)
        # Use image modality for latlon, elevation, other base data?
        image_modality = InputModality(
            name=TOPOGRAPHIC_DATA,
            input_channels=image_input_channels,
            input_axis=2,  # number of axes, 2 for images
            num_freq_bands=input_size,  # number of freq bands, with original value (2 * K + 1)
            max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is
            sin_only=sin_only,
            fourier_encode=encode_fourier,
        )
        modalities.append(image_modality)
        if not self.predict_timesteps_together:
            # Sort audio for timestep one-hot encode? Or include under other modality?
            timestep_modality = InputModality(
                name="forecast_time",
                input_channels=1,  # number of channels for mono audio
                input_axis=1,  # number of axes, 2 for images
                num_freq_bands=self.forecast_steps,  # number of freq bands, with original value (2 * K + 1)
                max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is
                sin_only=sin_only,
                fourier_encode=encode_fourier,
            )
            modalities.append(timestep_modality)
        # X,Y Coords are given in 1D, and each would be a different modality
        # Keeping them as 1D saves input size, just need to add more ones
        coord_modalities = (
            [
                SATELLITE_Y_COORDS,
                SATELLITE_X_COORDS,
                TOPOGRAPHIC_Y_COORDS,
                TOPOGRAPHIC_X_COORDS,
                NWP_Y_COORDS,
                NWP_X_COORDS,
            ]
            if nwp_modality
            else [
                SATELLITE_Y_COORDS,
                SATELLITE_X_COORDS,
                TOPOGRAPHIC_Y_COORDS,
                TOPOGRAPHIC_X_COORDS,
            ]
        )
        for coord in coord_modalities:
            coord_modality = InputModality(
                name=coord,
                input_channels=1,  # number of channels for mono audio
                input_axis=1,  # number of axes, 2 for images
                num_freq_bands=input_size,  # number of freq bands, with original value (2 * K + 1)
                max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is
                sin_only=sin_only,
                fourier_encode=encode_fourier,
            )
            modalities.append(coord_modality)

        # Datetime features as well should be incorporated
        if datetime_modality:
            for date in [SATELLITE_DATETIME_INDEX] + list(DATETIME_FEATURE_NAMES):
                date_modality = InputModality(
                    name=date,
                    input_channels=1,  # number of channels for mono audio
                    input_axis=1,  # number of axes, 2 for images
                    num_freq_bands=(
                        2 * history_steps + 1
                    ),  # number of freq bands, with original value (2 * K + 1)
                    max_freq=max_frequency,  # maximum frequency, hyperparameter depending on how fine the data is
                    sin_only=sin_only,
                    fourier_encode=encode_fourier,
                )
                modalities.append(date_modality)

        self.model = MultiPerceiver(
            modalities=modalities,
            dim=dim,  # dimension of sequence to be encoded
            queries_dim=queries_dim,  # dimension of decoder queries
            logits_dim=logits_dim,  # dimension of final logits
            depth=depth,  # depth of net
            num_latents=num_latents,  # number of latents, or induced set points, or centroids. different papers giving it different names
            latent_dim=latent_dim,  # latent dimension
            cross_heads=cross_heads,  # number of heads for cross attention. paper said 1
            latent_heads=latent_heads,  # number of heads for latent self attention, 8
            cross_dim_head=cross_dim_heads,  # number of dimensions per cross attention head
            latent_dim_head=latent_dim_heads,  # number of dimensions per latent self attention head
            weight_tie_layers=weight_tie_layers,  # whether to weight tie layers (optional, as indicated in the diagram)
            # self_per_cross_attn=self_per_cross_attention,  # number of self attention blocks per cross attention
            sine_only=sin_only,
            fourier_encode_data=encode_fourier,
            output_shape=input_size,  # Shape of output to make the correct sized logits dim, needed so reshaping works
            decoder_ff=decoder_ff,  # Optional decoder FF
        )

        if postprocessor_type is not None:
            if postprocessor_type not in ("conv", "patches", "pixels", "conv1x1"):
                raise ValueError("Invalid postprocessor_type!")
            self.postprocessor = ImageDecoder(
                postprocess_type=postprocessor_type, output_channels=sat_channels, **decoder_kwargs
            )
        else:
            self.postprocessor = None

        self.save_hyperparameters()

    def encode_inputs(self, x: dict) -> Dict[str, torch.Tensor]:
        video_inputs = x[SATELLITE_DATA]
        nwp_inputs = x.get(NWP_DATA, [])
        base_inputs = x.get(
            TOPOGRAPHIC_DATA, []
        )  # Base maps should be the same for all timesteps in a sample

        # Run the preprocessors here when encoding the inputs
        if self.preprocessor is not None:
            # Expects Channel first
            video_inputs = self.preprocessor(video_inputs)
            base_inputs = self.preprocessor(base_inputs)
            if nwp_inputs:
                nwp_inputs = self.preprocessor(nwp_inputs)
        video_inputs = video_inputs.permute(0, 1, 3, 4, 2)  # Channel last
        if nwp_inputs:
            nwp_inputs = nwp_inputs.permute(0, 1, 3, 4, 2)  # Channel last
            x[NWP_DATA] = nwp_inputs
        base_inputs = base_inputs.permute(0, 2, 3, 1)  # Channel last
        logger.debug(f"Timeseries: {video_inputs.size()} Base: {base_inputs.size()}")
        x[SATELLITE_DATA] = video_inputs
        x[TOPOGRAPHIC_DATA] = base_inputs
        return x

    def add_timestep(self, batch_size: int, timestep: int = 1) -> torch.Tensor:
        times = (torch.eye(self.forecast_steps)[timestep]).unsqueeze(-1).unsqueeze(-1)
        ones = torch.ones(1, 1, 1)
        timestep_input = times * ones
        timestep_input = timestep_input.squeeze(-1)
        timestep_input = repeat(timestep_input, "... -> b ...", b=batch_size)
        logger.debug(f"Forecast Step: {timestep_input.size()}")
        return timestep_input

    def _train_or_validate_step(self, batch, batch_idx, is_training: bool = True):
        x, y = batch
        batch_size = y[SATELLITE_DATA].size(0)
        # For each future timestep:
        predictions = []
        query = self.construct_query(x)
        x = self.encode_inputs(x)
        if self.predict_timesteps_together:
            # Predicting all future ones at once
            y_hat = self(x, query=query)
            y_hat = rearrange(
                y_hat,
                "b (t h w) c -> b c t h w",
                t=self.forecast_steps,
                h=self.input_size,
                w=self.input_size,
            )
        else:
            for i in range(self.forecast_steps):
                x["forecast_time"] = self.add_timestep(batch_size, i).type_as(y)
                y_hat = self(x, query=query)
                y_hat = rearrange(y_hat, "b h (w c) -> b c h w", c=self.output_channels)
                predictions.append(y_hat)
            y_hat = torch.stack(predictions, dim=1)  # Stack along the timestep dimension
        if self.postprocessor is not None:
            y_hat = self.postprocessor(y_hat)
        if self.visualize:
            self.visualize_step(x, y, y_hat, batch_idx, step="train" if is_training else "val")
        loss = self.criterion(y, y_hat)
        self.log_dict({f"{'train' if is_training else 'val'}/loss": loss})
        frame_loss_dict = {}
        for f in range(self.forecast_steps):
            frame_loss = self.criterion(
                y_hat[:, f, :, :, :], y[SATELLITE_DATA][:, f, :, :, :]
            ).item()
            frame_loss_dict[f"{'train' if is_training else 'val'}/frame_{f}_loss"] = frame_loss
        self.log_dict(frame_loss_dict)
        return loss

    def configure_optimizers(self):
        # They use LAMB as the optimizer
        optimizer = optim.Lamb(self.parameters(), lr=self.lr, betas=(0.9, 0.999))
        scheduler = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=10, max_epochs=100)
        lr_dict = {
            # REQUIRED: The scheduler instance
            "scheduler": scheduler,
            # The unit of the scheduler's step size, could also be 'step'.
            # 'epoch' updates the scheduler on epoch end whereas 'step'
            # updates it after a optimizer update.
            "interval": "step",
            # How many epochs/steps should pass between calls to
            # `scheduler.step()`. 1 corresponds to updating the learning
            # rate after every epoch/step.
            "frequency": 1,
            # If using the `LearningRateMonitor` callback to monitor the
            # learning rate progress, this keyword can be used to specify
            # a custom logged name
            "name": None,
        }
        return {"optimizer": optimizer, "lr_scheduler": lr_dict}

    def construct_query(self, x: dict):
        if self.use_learnable_query:
            if self.temporally_consistent_fourier_features:
                fourier_features = encode_position(
                    x[SATELLITE_DATA].shape[0],
                    axis=(
                        x[SATELLITE_DATA].shape[1] + self.forecast_steps,
                        self.input_size,
                        self.input_size,
                    ),
                    num_frequency_bands=max(
                        [self.input_size, x[SATELLITE_DATA].shape[1] + self.forecast_steps]
                    )
                    * 2
                    + 1,
                    max_frequency=self.max_frequency,
                )[
                    x[SATELLITE_DATA].shape[1] :
                ]  # Only want future part
            else:
                fourier_features = None
            return self.query(x, fourier_features)
        # key, value: B x N x K; query: B x M x K
        # Attention maps -> B x N x M
        # Output -> B x M x K
        # So want query to be B X (T*H*W) X C to reshape to B x T x C x H x W
        if self.preprocessor is not None:
            x = self.preprocessor(x[SATELLITE_DATA])
        y_query = x  # Only want sat channels, the output
        # y_query = torch.permute(y_query, (0, 2, 3, 1)) # Channel Last
        # Need to reshape to 3 dimensions, TxHxW or HxWxC
        # y_query = rearrange(y_query, "b h w d -> b (h w) d")
        logger.debug(f"Query Shape: {y_query.shape}")
        return y_query

    def forward(self, x, mask=None, query=None):
        return self.model.forward(x, mask=mask, queries=query)


# scrape data into CSV
# grab all links from rithm school blog
# store url, anchor tag text and date

import requests
from bs4 import BeautifulSoup
from csv import writer

def web_crawl(url):
    response = requests.get(url)

    soup = BeautifulSoup(response.text, "html.parser")
    articles = soup.find_all('article')

    with open('web_scraping/blog_data.csv', 'a') as file:
        csv_writer = writer(file)
        csv_writer.writerow(["title","link","date"])

        for article in articles:
            a_tag = article.find('a')
            title = a_tag.get_text()
            a_url = a_tag['href']
            date = article.find('time')['datetime']
            csv_writer.writerow([title, a_url, date])

    try:
        next_page = soup.find(class_ = 'current').find_next_sibling().find('a')['href']
    except:
        exit()

    next_page = next_page[5:]

    print(next_page)

    if url[-1].isdigit():
        next_url = url[:-1] + next_page[-1]
    else:
        next_url = url + next_page

    web_crawl(next_url)

    return

web_crawl('https://www.rithmschool.com/blog')


from django.urls import path, include, re_path
from rest_framework.routers import DefaultRouter

from post.views import ProfileFeed
from . import views
from question.views import ProfilePageViewSet, QuestionViewSet, CommentViewSet

router = DefaultRouter()

router.register(
    r"profile/(?P<accounts_pk>\d+)/pages", ProfilePageViewSet, basename="profile-pages",
)
router.register(
    r"profile/(?P<accounts_pk>\d+)/pages/(?P<page_pk>\d+)/questions",
    QuestionViewSet,
    basename="profile-questions",
)
router.register(
    r"profile/(?P<accounts_pk>\d+)/pages/(?P<page_pk>\d+)/questions/(?P<question_pk>\d+)/comments",
    CommentViewSet,
    basename="profile-comment",
)
router.register(
    r"profile/(?P<accounts_pk>\d+)/feed", ProfileFeed, basename="profile-cards",
)

urlpatterns = [
    path("profile/<int:pk>/", views.ProfilePage.as_view()),
    re_path(
        r"profile/(?P<accounts_pk>\d+)/mento/", views.mento_user, name="user-mento"
    ),
    re_path(
        r"profile/(?P<accounts_pk>\d+)/mentiee/",
        views.mentiee_user,
        name="user-mentiee",
    ),
    path("", include(router.urls)),
    path("accounts/signup/", views.signup_view, name="signup"),
    path("accounts/login/", views.login_view, name="login"),
    path("accounts/tendency/", views.tendency_view, name="tendency"),
    path("accounts/tmp-password/", views.tmp_password, name="tmp-password"),
    path("accounts/delete/", views.DeleteUserView.as_view(), name="delete-user"),
    path(
        "accounts/forgot-username/", views.forgot_username_view, name="forgot-username"
    ),
    path(
        "accounts/comment/",
        views.UpdateDecsription.as_view(),
        name="update-description",
    ),
    path("accounts/password/", views.ChangePassword.as_view(), name="update-password",),
    path("accounts/username/", views.ChangeUsername.as_view(), name="update-username",),
    path("accounts/image/", views.UpdateUserImage.as_view(), name="update-image",),
]


from jax import numpy as jnp
from jax.scipy.spatial.transform import Rotation

class CameraConfig:
    def __init__(self, **kwargs) -> None:
        self.url = kwargs.pop("url", "")

        self.calibration = jnp.array(
            kwargs.pop(
                "calibration",
                [
                    [100, 0, 50],
                    [0, 100, 50],
                    [0, 0, 1],
                ],
            )
        )

        distortion = kwargs.pop("distortion", None)
        if distortion is not None:
            self.distortion = jnp.array(*distortion)
        else:
            self.distortion = None

        self.position = jnp.array(kwargs.pop("position"))
        orientation: Rotation | list[float] = kwargs.pop("orientation")

        if isinstance(orientation, list):
            self.orientation = Rotation.from_euler(
                "ZYX", jnp.array(orientation), degrees=True
            )
        else:
            self.orientation = orientation


import matplotlib.pyplot as plt
import pandas as pd
import os
import seaborn as sns
plt.ion()

data = pd.DataFrame()
for x in os.listdir("./output"):
    p = pd.read_excel("./output/"+x)
    data = data.append(p, ignore_index=True)

data = data.loc[data.amplitude < data.rms * -1.645]
data = data.loc[data.amplitude > -500]
data = data.loc[data.tau > 0.001]
data.tau_pams = data.tau_pams / 1000
data.rise_pams = data.rise_pams / 1000

wt = [30718, 130618, 190718, 250718]
data["geno"] = ["wt" if x in wt else "mutant" for x in data["date"]]

fig, ax = plt.subplots(6,4)
for i, x in enumerate(["amplitude", "rise", "tau", "rise_pams", "tau_pams"]):
    p = data.groupby(["date", "cell"]).mean().reset_index()
    sns.barplot(x="date", y=x, data=p, ax=ax[i,0], alpha=.5, ci=False)
    sns.stripplot(x="date", y=x, data=p, ax=ax[i,0], jitter=True, hue="cell")
    for y in data["date"].unique():
        sns.kdeplot(data.loc[data.date == y, x], ax=ax[i, 1], cumulative=True, label=y)
p = data.groupby(["date", "cell"]).count().reset_index()
sns.barplot(x="date", y="rms", data=p, alpha=.5, ci=False, ax=ax[5,0])
sns.stripplot(x="date", y="rms", data=p, jitter=True, ax=ax[5,0])
ax[5,0].set_ylabel("Event Count")

for i, x in enumerate(["amplitude", "rise", "tau", "rise_pams", "tau_pams"]):
    p = data.groupby(["date", "cell", "geno"]).mean().reset_index()
    sns.barplot(x="geno", y=x, data=p, ax=ax[i,2], alpha=.5, ci=False)
    sns.stripplot(x="geno", y=x, data=p, ax=ax[i,2], jitter=True, hue="cell")
    for y in data["geno"].unique():
        sns.kdeplot(data.loc[data.geno == y, x], ax=ax[i, 3], cumulative=True, label=y)
p = data.groupby(["date", "cell", "geno"]).count().reset_index()
sns.barplot(x="geno", y="rms", data=p, alpha=.5, ci=False, ax=ax[5,2])
sns.stripplot(x="geno", y="rms", data=p, jitter=True, ax=ax[5,2])
ax[5,0].set_ylabel("Event Count")

for a in ax.flatten():
    a.legend_.remove()

plt.tight_layout()


import os
import sys
import cv2
import time
import pickle
import numpy as np
from scipy import io
import matplotlib.pyplot as plt

from config import cfg
from facial_shape import recon_facial_shape
from landmark import detect_lmk_2d, detect_lmk_3d
from helpers import relative, head_pose, eye_mask, pupil_center, eyeball_center, average_gaze, lines_intersection, W_2_I, W_2_C, I_2_W
from visualization import visualize_3d

sys.path.append('./TDDFA_V2/')

if __name__ == '__main__':
    subj_idx = cfg.EasyCali.subj_idx

    image_save_dir = os.path.join('./calibration/', cfg.EasyCali.image_save_dir, 'subject_{}'.format(subj_idx))
    image_name_lst = os.listdir(image_save_dir)
    if 'intrinsic_params.mat' in image_name_lst:
        image_name_lst.remove('intrinsic_params.mat')
    image_name_lst.sort()
    print('Now:', len(image_name_lst))

    intrinsic_params = io.loadmat(os.path.join(image_save_dir, 'intrinsic_params.mat'))

    with open(cfg.lmk_idx_path, 'rb') as lmk_idx_file:
        lmk_idx_dict = pickle.load(lmk_idx_file)

    result_save_path = os.path.join('./calibration/', cfg.EasyCali.result_save_dir, 'subject_{}'.format(subj_idx))
    if not os.path.exists(result_save_path):
        os.makedirs(result_save_path)

    # Facial shape recovery
    fs_3d_path = os.path.join(result_save_path, 'facial_shape_3d.npy')
    lmk_3d_path = os.path.join(result_save_path, 'landmark_3d.npy')

    if not os.path.exists(fs_3d_path) or not os.path.exists(lmk_3d_path):
        assert cfg.face_3d_model in ['3DDFA_V2']
        all_fs_3d = np.zeros((len(image_name_lst), cfg.face_3d_model_shape[0], 3))

        for image_idx in range(len(image_name_lst)):
            image_name = image_name_lst[image_idx]
            image_path = os.path.join(image_save_dir, image_name)

            image = cv2.imread(image_path)
            fs_3d, fs_frontal_3d, exist_fs = recon_facial_shape(image, cfg)

            all_fs_3d[image_idx, :, :] = fs_3d
            print(image_name, '-- facial shape done.')

        avg_fs_3d = np.mean(all_fs_3d, axis=0)
        print(avg_fs_3d.shape)
        np.save(fs_3d_path, avg_fs_3d)

        lmk_3d = detect_lmk_3d(avg_fs_3d, lmk_idx_dict)
        np.save(lmk_3d_path, lmk_3d)

    else:
        avg_fs_3d = np.load(fs_3d_path)
        lmk_3d = np.load(lmk_3d_path)

    fs_3d, fs_frontal_3d = avg_fs_3d, avg_fs_3d[: cfg.face_3d_model_frontal_idx, :]

    # Eyeball center calibration
    EyeballCenter_path = os.path.join(result_save_path, 'eyeball_center.npy')
    gaze_lines_path = os.path.join(result_save_path, 'gaze_lines.mat')

    if not os.path.exists(EyeballCenter_path) or not os.path.exists(gaze_lines_path):

        left_lines_strt = np.zeros((len(image_name_lst), 3))
        left_lines_vec = np.zeros((len(image_name_lst), 3))
        right_lines_strt = np.zeros((len(image_name_lst), 3))
        right_lines_vec = np.zeros((len(image_name_lst), 3))

        no_lmk_idx = []

        for image_idx in range(len(image_name_lst)):
            image_name = image_name_lst[image_idx]
            image_path = os.path.join(image_save_dir, image_name)

            image = cv2.imread(image_path)

            # 2d landmark detection
            assert cfg.face_2d_model in ['MediaPipe']
            lmk_2d, Pupil_2d, Eye_contour_2d, exist_lmk = detect_lmk_2d(image, lmk_idx_dict, cfg)

            if not exist_lmk:
                no_lmk_idx.append(image_idx)
                continue

            # Head pose
            rot_v, rot_mat, trans_v = head_pose(lmk_2d, lmk_3d, intrinsic_params)

            # W ---> I
            fs_W, fs_frontal_W = fs_3d, fs_frontal_3d
            fs_frontal_I = W_2_I(fs_frontal_W, rot_v, trans_v, intrinsic_params)

            # Two points for one gaze line
            Pupil_W_1 = I_2_W(Pupil_2d, rot_mat, trans_v, intrinsic_params, scalar=3e5)
            Pupil_W_2 = I_2_W(Pupil_2d, rot_mat, trans_v, intrinsic_params, scalar=6e5)

            left_lines_strt[image_idx, :] = Pupil_W_1[0]
            left_lines_vec[image_idx, :] = Pupil_W_2[0] - Pupil_W_1[0]
            right_lines_strt[image_idx, :] = Pupil_W_1[1]
            right_lines_vec[image_idx, :] = Pupil_W_2[1] - Pupil_W_1[1]

            print(image_name, '-- gaze line done.')

        left_lines_strt = np.delete(left_lines_strt, no_lmk_idx, axis=0)
        left_lines_vec = np.delete(left_lines_vec, no_lmk_idx, axis=0)
        right_lines_strt = np.delete(right_lines_strt, no_lmk_idx, axis=0)
        right_lines_vec = np.delete(right_lines_vec, no_lmk_idx, axis=0)

        LeftEyeballCenter_W = lines_intersection(left_lines_strt, left_lines_vec)
        RightEyeballCenter_W = lines_intersection(right_lines_strt, right_lines_vec)

        EyeballCenter_W = np.array([LeftEyeballCenter_W, RightEyeballCenter_W])
        np.save(EyeballCenter_path, EyeballCenter_W)

        gaze_lines_dict = {
            'left_lines_strt': left_lines_strt,
            'left_lines_vec': left_lines_vec,
            'right_lines_strt': right_lines_strt,
            'right_lines_vec': right_lines_vec,
        }
        io.savemat(gaze_lines_path, mdict=gaze_lines_dict)

    else:
        EyeballCenter_W = np.load(EyeballCenter_path)
        LeftEyeballCenter_W, RightEyeballCenter_W = EyeballCenter_W[0], EyeballCenter_W[1]

        gaze_lines_dict = io.loadmat(gaze_lines_path)

    # Visualize
    if cfg.EasyCali.visualize:
        vis_3d_dict_W = {  # WCS default
            'CS': 'W',
            'view': (90, -90),

            'facial_shape_3d': {'obj': fs_3d, 's': 20, 'alpha': 0.05},
            'LeftEyeballCenter': {'obj': LeftEyeballCenter_W, 'color': 'g', 's': 40, 'alpha': 0.8},
            'RightEyeballCenter': {'obj': RightEyeballCenter_W, 'color': 'g', 's': 40, 'alpha': 0.8},
            'landmark': {'obj': lmk_3d, 'color': 'r', 's': 40, 'alpha': 0.8},
            'gaze_lines_dict': {'obj': gaze_lines_dict, 'sampling_rate': 1/5, 'start_pos': 0.3, 'end_pos': 0.7,
                                'left_color': 'r', 'right_color': 'b', 'linewidth': 2.0, 'alpha': 0.4},

            'show_tick': True,  # [True, False]
        }

        visualize_3d(vis_3d_dict_W)


#!/usr/bin/python3

from datetime import datetime as dt
from datetime import timedelta
from pytz import timezone as tz
import pytz
from operator import itemgetter
from itertools import chain

class AirPort:
    def __init__(self, port_id, country_id,timezone):
        self.code = port_id
        self.country = country_id
        self.tzone = timezone
        self.dstntions = []
    def add_destination(self, *args):
        arr_port, dep_time, arr_time = args
        dep_zone = tz(self.tzone)
        arr_zone = tz(airports[arr_port].tzone)
        dep_time = dep_zone.localize(dt.strptime(dep_time, "%Y-%m-%d %H:%M:%S"))
        arr_time = arr_zone.localize(dt.strptime(arr_time, "%Y-%m-%d %H:%M:%S"))
        self.dstntions.append((arr_port,dep_time,arr_time))

####################################################
airports = {}
with open('iata_codes.csv','r') as iatafile:
    for line in iatafile:
        ia_country, ia_port, ia_timezone = line.rstrip().split(',')
        airports[ia_port]=AirPort(ia_port,ia_country,ia_timezone)

# parse input data
with open('input_data.csv','r') as datafile:
    for line in datafile:
        dep_port, arr_port, dep_time, arr_time = line.rstrip().split(';')
        #cut off destination in the same country
        if airports[dep_port].country != airports[arr_port].country:
            airports[dep_port].add_destination(arr_port,dep_time,arr_time)
####################################################

def date_to_string(ptimes):
    t_fmt='%Y-%m-%dT%H:%M:%S'
    return [(t[0].strftime(t_fmt), t[1].strftime(t_fmt)) for t in ptimes]

def checktime(nxt,pathtime):
    if not pathtime or  (nxt[1]>pathtime[-1][1] and nxt[2]-pathtime[0][0]<path_maxlen):
        return True,nxt[1:]
    return False, None

def enumpaths(k,c,visitedcntrs=[],path=[],pathtime=[]):
    if len(pathlist)>40*c:
        raise ValueError('Limit of 40 paths is reached, choosing next start port..')
    palen = len(path)
    if not visitedcntrs:
        visitedcntrs=[airports[k].country]
    if palen>9:
        pathlist.append(path)
        timelist.append(date_to_string(pathtime))
    else:
        for nxt in airports[k].dstntions:
            cun = airports[nxt[0]].country
            t_ok, t_tup = checktime(nxt,pathtime)
            if t_ok and ((palen==9 and cun==visitedcntrs[0]) or (palen<9 and cun not in visitedcntrs)):
                enumpaths(nxt[0],c,visitedcntrs+[cun],path+[(visitedcntrs[-1],k,nxt[0])],pathtime+[t_tup])

def formatresults():
    rez=[list(zip(pathlist[i],timelist[i])) for i in range(len(pathlist))]
    rez=[[str(i)]+list(chain.from_iterable(y)) for i,x in enumerate(rez) for y in x]
    return rez

def export_res(rez):
    with open('results_10.csv','a') as rzf:
        for line in rez:
            rzf.write(';'.join(line)+'\n')

pathlist, timelist = [], []
path_maxlen=timedelta(365)
c=0
for port in airports:
    c += 1
    try:
        enumpaths(port,c)
    except ValueError as e:
        print(port,e)
    if c>40:
        break

rez=formatresults()
export_res(rez)


# coding=utf-8
import time
import os
import json

import pandas as pd
import requests

from src import urls_getter
from src import page_parser

root_url = "http://sofasofa.io/"  # 主站地址
page_interface = "http://sofasofa.io/public_forum_exe.php?action=load_more"  # 分页请求接口

url_pools = './data/urls_pool.csv'
url_pools_df = None
if os.path.exists(url_pools):
    url_pools_df = pd.read_csv(url_pools, index_col=[0])
else:
    with open(url_pools, 'a+') as fw:
        for i in range(171):
            print('*' * 100)
            print(f'get {i}th page')
            payload = {"start": i*10, "filter": None, "type": 0}  # 接口传递参数
            get_url = urls_getter.UrlGetter()
            get_url.get_entities(root_url, page_interface, 'POST', payload)
            titles = get_url.get_urls()
            fw.writelines('\n'.join(titles))
            fw.write('\n')
            time.sleep(0.1)

    url_pools_df = pd.read_csv(url_pools, header=None)
    url_pools_df.columns = ['url']
    url_pools_df = url_pools_df.sort_values(by='url').reset_index(drop=True)
    url_pools_df.to_csv(url_pools)

titles = url_pools_df['url'].to_list()
parse_page = page_parser.Parser()
i = 0
for question_url in titles:
    print('*'*100)
    print(f'get page of {question_url}')
    id = question_url.split('=')[1]
    parse_page.get_page(question_url)
    title = parse_page.get_question()
    print(f'questions: {title}')
    desc = parse_page.get_question_desc()
    viewer =parse_page.get_viewer()
    tags = parse_page.get_tags()
    answers = parse_page.get_answer()
    print(answers)
    QA = {
        'id': id,
        'question': title,
        'description': desc,
        'viewer': viewer,
        'tags': tags,
        'answers': answers
    }
    with open('./data/output/' + id + '.json', 'w') as fw:
        json.dump(QA, fw, ensure_ascii=False)


import os
from dotenv import load_dotenv
import time
import random
import wandb

from UNet import RecursiveUNet
from engine import Segmentor
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts,CosineAnnealingLR, ReduceLROnPlateau
from dataloader import prep_dataloader

import numpy as np
from torch import nn
from torch.cuda import amp
import torch
import SimpleITK as sitk
sitk.ProcessObject_SetGlobalWarningDisplay(False)

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

def seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

class CFG:

    def __init__(self):
        self.data_path = os.getenv('VIDA_PATH')
        self.name = "UNet_baseline"
        self.model = 'RecursiveUNet' # ['UNet', 'RecursiveUNet']
        self.optimizer = 'adam' # ['adam', 'adamp']
        self.scheduler = 'CosineAnnealingWarmRestarts'
        # ['CosineAnnealingWarmRestarts', 'CosineAnnealingLR', 'ReduceLROnPlateau']
        self.epochs = 10
        self.T_0 = self.epochs

        self.in_file = 'ENV18PM_ProjSubjList_cleaned_IN.in'
        self.lr = 0.0001
        self.min_lr = 1e-6
        self.train_bs = 16
        self.valid_bs = 32
        self.test_results_dir = "RESULTS"
        self.mask = 'airway'
        self.save = False
        self.debug = False

def tune_hyper():

    with wandb.init() as run:
        print('config:', dict(run.config))
        # Model
        if run.config.activation == 'elu':
            activation = nn.ELU(inplace=True)
        else:
            activation = nn.ReLU(inplace=True)

        model = RecursiveUNet(activation=activation)
        # model = RecursiveUNet(activation=nn.ELU(inplace=True))
        DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model.to(DEVICE)
        loss_function = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(),lr=0.0001)
        scheduler = CosineAnnealingWarmRestarts(optimizer,
                                                    T_0=c.T_0,
                                                    T_mult=1,
                                                    eta_min=1e-6,
                                                    last_epoch=-1)

        scaler = amp.GradScaler()
        eng = Segmentor(model=model,
                    optimizer=optimizer,
                    scheduler=scheduler,
                    loss_fn=loss_function,
                    device=DEVICE,
                    scaler=scaler)

        best_loss = np.inf
        best_dice = 0
        # Train
        wandb.watch(eng.model,log='all',log_freq=10)
        for epoch in range(c.epochs):
            trn_loss, trn_dice = eng.train(train_loader)
            val_loss, val_dice = eng.evaluate(valid_loader)
            wandb.log({'epoch': epoch, 'trn_loss': trn_loss, 'val_loss': val_loss,
                        'trn_dice': trn_dice, 'val_dice': val_dice})
            if c.scheduler == 'ReduceLROnPlateau':
                scheduler.step(val_loss)
            eng.epoch += 1
            print(f'Epoch: {epoch}, train loss: {trn_loss:5f}, valid loss: {val_loss:5f}')
            print(f'Epoch: {epoch}, train dice: {trn_dice:5f}, valid dice: {val_dice:5f}')
        #  if val_dice > best_dice:
        #      best_dice = val_dice
        #      print('Best Dice:', best_dice)
        #      if c.save:
        #          torch.save(model.state_dict(), path)
            if val_loss < best_loss:
                best_loss = val_loss
                best_dice = val_dice
                print(f'Best loss: {best_loss} with {best_dice}')
        torch.cuda.empty_cache()

if __name__ == "__main__":
    load_dotenv()
    seed_everything()
    # wandb.init(project='airway')

    sweep_config = {
        'name':'UNet',
        'method':'bayes',
        'metric': {
            'name': 'val_loss',
            'goal': 'minimize'
        },
        'parameters': {
            'activation': {
                'values' : ['elu','relu']
            }
        }
        # 'parameters': {
        #     'epochs': {
        #         'values' : [10,20]
        #     },
        #     'learning_rate' : {
        #         'min' : 0.0001,
        #         'max' : 0.01
        #     }
        # }

    }
    sweep_id = wandb.sweep(sweep_config,project='airway')

    c = CFG()
    n_case = 64
    print('***********************************************************')
    print('Configuration: ')
    print(c.__dict__)
    print('***********************************************************')

    # Data
    if c.debug: # only use 10 cases, 1 epoch
        train_loader, valid_loader = prep_dataloader(c,n_case=5)
        c.epochs = 2
    else:
        train_loader, valid_loader = prep_dataloader(c,n_case=n_case)

    count = 2
    wandb.agent(sweep_id, function=tune_hyper, count=count)


# add your save-note function here

import boto3
import json

dynamob_resource = boto3.resource("dynamodb")

table = dynamob_resource.Table("lotion-30112955")

'''def put_item(note):
    table.put_item(
        Item = note
    )
    return*'''

def handler(event, context):

    body = json.loads(event['body'])
    try:
        table.put_item(Item = body)
        return{
            'statusCode':200,
            "body" : json.dumps({
            "message": "Note added successfully"
            }
            )
        }
    except Exception as e :
        print(f'exception:{e}')
        return{
            'statusCode':500,
            "body": json.dumps({
            "message": str(e)
            })
        }