import tensorflow as tf

from tensorforce import TensorforceError, util
from tensorforce.core import tf_util
from tensorforce.core.parameters import Parameter

class Decaying(Parameter):
    """
    Decaying hyperparameter (specification key: `decaying`, `linear`, `exponential`, `polynomial`,
    `inverse_time`, `cosine`, `cosine_restarts`, `linear_cosine`, `linear_cosine_noisy`).

    Args:
        decay ("linear" | "exponential" | "polynomial" | "inverse_time" | "cosine" | "cosine_restarts" | "linear_cosine" | "linear_cosine_noisy"):
            Decay type, see also
            `TensorFlow docs <http://example.com>`__
            (required).
        unit ("timesteps" | "episodes" | "updates"): Unit of decay schedule
            (required).
        num_steps (int): Number of decay steps
            (required).
        initial_value (float | int): Initial value
            (required).
        increasing (bool): Whether to subtract the decayed value from 1.0
            (default: false).
        inverse (bool): Whether to take the inverse of the decayed value
            (default: false).
        scale (float): Scaling factor for (inverse) decayed value
            (default: 1.0).
        kwargs: Additional arguments depend on decay mechanism.

            Linear decay:

final_value (float | int) – Final value
            (required).

            Exponential decay:

decay_rate (float) – Decay rate
            (required).

staircase (bool) – Whether to apply decay in a discrete
            staircase, as opposed to continuous, fashion.
            (default: false).

            Polynomial decay:

final_value (float | int) – Final value
            (required).

power (float | int) – Power of polynomial
            (default: 1, thus linear).

cycle (bool) – Whether to cycle beyond num_steps
            (default: false).

            Inverse time decay:

decay_rate (float) – Decay rate
            (required).

staircase (bool) – Whether to apply decay in a discrete
            staircase, as opposed to continuous, fashion.
            (default: false).

            Cosine decay:

alpha (float) – Minimum learning rate value as a fraction of
            learning_rate
            (default: 0.0).

            Cosine decay with restarts:

t_mul (float) – Used to derive the number of iterations in the
            i-th period
            (default: 2.0).

m_mul (float) – Used to derive the initial learning rate of the
            i-th period
            (default: 1.0).

alpha (float) – Minimum learning rate value as a fraction of
            the learning_rate
            (default: 0.0).

            Linear cosine decay:

num_periods (float) – Number of periods in the cosine part of
            the decay
            (default: 0.5).

alpha (float) – Alpha value
            (default: 0.0).

beta (float) – Beta value
            (default: 0.001).

            Noisy linear cosine decay:

initial_variance (float) – Initial variance for the noise
            (default: 1.0).

variance_decay (float) – Decay for the noise's variance
            (default: 0.55).

num_periods (float) – Number of periods in the cosine part of
            the decay
            (default: 0.5).

alpha (float) – Alpha value
            (default: 0.0).

beta (float) – Beta value
            (default: 0.001).

        name (string): internal use.
        dtype (type): internal use.
        min_value (dtype-compatible value): internal use.
        max_value (dtype-compatible value): internal use.
    """

    def __init__(
        self, *, decay, unit, num_steps, initial_value, increasing=False, inverse=False, scale=1.0,
        name=None, dtype=None, min_value=None, max_value=None, **kwargs
    ):
        assert decay in (
            'cosine', 'cosine_restarts', 'exponential', 'inverse_time', 'linear', 'linear_cosine',
            'linear_cosine_noisy', 'polynomial'
        )
        assert unit in ('timesteps', 'episodes', 'updates')
        assert (isinstance(num_steps, int) or num_steps % 10.0 == 0.0) and num_steps > 0
        assert isinstance(initial_value, (float, int))

        if isinstance(initial_value, int):
            if dtype != 'int':
                raise TensorforceError.dtype(
                    name='Decaying', argument='initial_value', dtype=type(initial_value)
                )
        elif isinstance(initial_value, float):
            if dtype != 'float':
                raise TensorforceError.dtype(
                    name='Decaying', argument='initial_value', dtype=type(initial_value)
                )
        else:
            raise TensorforceError.unexpected()

        if decay == 'linear':
            assert len(kwargs) == 1 and 'final_value' in kwargs
            decay = 'polynomial'

        self.decay = decay
        self.num_steps = int(num_steps)
        self.initial_value = initial_value
        self.increasing = increasing
        self.inverse = inverse
        self.scale = scale
        self.kwargs = kwargs

        super().__init__(
            unit=unit, name=name, dtype=dtype, min_value=min_value, max_value=max_value
        )

    def min_value(self):
        if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('alpha', 0.0)
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= 0.0:
                min_value = 0.0
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('beta', 0.001)
                max_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
            else:
                min_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
                max_value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= self.kwargs['final_value']:
                min_value = self.kwargs['final_value']
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.kwargs['final_value']

        assert min_value <= max_value

        if self.increasing:
            assert 0.0 <= min_value <= max_value <= 1.0
            min_value, max_value = 1.0 - max_value, 1.0 - min_value

        if self.inverse:
            assert util.epsilon <= min_value <= max_value
            min_value, max_value = 1.0 / max_value, 1.0 / min_value

        if self.scale == 1.0:
            pass
        elif self.scale >= 0.0:
            min_value, max_value = self.scale * min_value, self.scale * max_value
        else:
            min_value, max_value = self.scale * max_value, self.scale * min_value

        return self.spec.py_type()(min_value)

if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('alpha', 0.0)
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= 0.0:
                min_value = 0.0
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('beta', 0.001)
                max_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
            else:
                min_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
                max_value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= self.kwargs['final_value']:
                min_value = self.kwargs['final_value']
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.kwargs['final_value']

        assert min_value <= max_value

        if self.increasing:
            assert 0.0 <= min_value <= max_value <= 1.0
            min_value, max_value = 1.0 - max_value, 1.0 - min_value

        if self.inverse:
            assert 0.0 < min_value <= max_value
            min_value, max_value = 1.0 / max_value, 1.0 / min_value

        if self.scale == 1.0:
            pass
        elif self.scale >= 0.0:
            min_value, max_value = self.scale * min_value, self.scale * max_value
        else:
            min_value, max_value = self.scale * max_value, self.scale * min_value

        return self.spec.py_type()(max_value)

    def final_value(self):
        if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                value = self.initial_value
            else:
                value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                value = self.initial_value
            else:
                value = self.kwargs['final_value']

        if self.increasing:
            assert 0.0 <= value <= 1.0
            value = 1.0 - value

        if self.inverse:
            assert value > 0.0
            value = 1.0 / value

        if self.scale != 1.0:
            value = value * self.scale

        return self.spec.py_type()(value)

    def parameter_value(self, *, step):
        initial_value = tf_util.constant(value=self.initial_value, dtype='float')

        if self.decay == 'cosine':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            parameter = tf.keras.experimental.CosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                alpha=self.kwargs.get('alpha', 0.0)
            )(step=step)

        elif self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            parameter = tf.keras.experimental.CosineDecayRestarts(
                initial_learning_rate=initial_value, first_decay_steps=(self.num_steps + 1),
                t_mul=self.kwargs.get('t_mul', 2.0), m_mul=self.kwargs.get('m_mul', 1.0),
                alpha=self.kwargs.get('alpha', 0.0)
            )(step=step)

        elif self.decay == 'exponential':
            assert self.kwargs['decay_rate'] >= 0.0
            parameter = tf.keras.optimizers.schedules.ExponentialDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                decay_rate=self.kwargs['decay_rate'], staircase=self.kwargs.get('staircase', False)
            )(step=step)

        elif self.decay == 'inverse_time':
            assert self.kwargs['decay_rate'] >= 0.0
            parameter = tf.keras.optimizers.schedules.InverseTimeDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                decay_rate=self.kwargs['decay_rate'], staircase=self.kwargs.get('staircase', False)
            )(step=step)

        elif self.decay == 'linear_cosine':
            assert self.kwargs.get('beta', 0.001) >= 0.0
            parameter = tf.keras.experimental.LinearCosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                num_periods=self.kwargs.get('num_periods', 0.5),
                alpha=self.kwargs.get('alpha', 0.0), beta=self.kwargs.get('beta', 0.001)
            )(step=step)

        elif self.decay == 'linear_cosine_noisy':
            assert self.kwargs.get('beta', 0.001) >= 0.0
            parameter = tf.keras.experimental.NoisyLinearCosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                initial_variance=self.kwargs.get('initial_variance', 1.0),
                variance_decay=self.kwargs.get('variance_decay', 0.55),
                num_periods=self.kwargs.get('num_periods', 0.5),
                alpha=self.kwargs.get('alpha', 0.0), beta=self.kwargs.get('beta', 0.001)
            )(step=step)

        elif self.decay == 'polynomial':
            assert self.kwargs.get('power', 1.0) >= 0.0
            parameter = tf.keras.optimizers.schedules.PolynomialDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                end_learning_rate=self.kwargs['final_value'], power=self.kwargs.get('power', 1.0),
                cycle=self.kwargs.get('cycle', False)
            )(step=step)

        if self.increasing:
            one = tf_util.constant(value=1.0, dtype='float')
            assertions = list()
            if self.config.create_tf_assertions:
                zero = tf_util.constant(value=0.0, dtype='float')
                assertions.append(tf.debugging.assert_greater_equal(x=parameter, y=zero))
                assertions.append(tf.debugging.assert_less_equal(x=parameter, y=one))
            with tf.control_dependencies(control_inputs=assertions):
                parameter = one - parameter

        if self.inverse:
            zero = tf_util.constant(value=0.0, dtype='float')
            epsilon = tf_util.constant(value=util.epsilon, dtype='float')
            parameter = tf.where(
                condition=(parameter > zero), x=(parameter + epsilon), y=(parameter - epsilon)
            )
            parameter = tf.math.reciprocal(x=parameter)

        if self.scale != 1.0:
            scale = tf_util.constant(value=self.scale, dtype='float')
            parameter = parameter * scale

        parameter = tf_util.cast(x=parameter, dtype=self.spec.type)

        return parameter


import tensorflow as tf

from tensorforce import TensorforceError, util
from tensorforce.core import tf_util
from tensorforce.core.parameters import Parameter

class Decaying(Parameter):
    """
    Decaying hyperparameter (specification key: `decaying`, `linear`, `exponential`, `polynomial`,
    `inverse_time`, `cosine`, `cosine_restarts`, `linear_cosine`, `linear_cosine_noisy`).

    Args:
        decay ("linear" | "exponential" | "polynomial" | "inverse_time" | "cosine" | "cosine_restarts" | "linear_cosine" | "linear_cosine_noisy"):
            Decay type, see also
            `TensorFlow docs <http://example.com>`__
            (required).
        unit ("timesteps" | "episodes" | "updates"): Unit of decay schedule
            (required).
        num_steps (int): Number of decay steps
            (required).
        initial_value (float | int): Initial value
            (required).
        increasing (bool): Whether to subtract the decayed value from 1.0
            (default: false).
        inverse (bool): Whether to take the inverse of the decayed value
            (default: false).
        scale (float): Scaling factor for (inverse) decayed value
            (default: 1.0).
        kwargs: Additional arguments depend on decay mechanism.

            Linear decay:

final_value (float | int) – Final value
            (required).

            Exponential decay:

decay_rate (float) – Decay rate
            (required).

staircase (bool) – Whether to apply decay in a discrete
            staircase, as opposed to continuous, fashion.
            (default: false).

            Polynomial decay:

final_value (float | int) – Final value
            (required).

power (float | int) – Power of polynomial
            (default: 1, thus linear).

cycle (bool) – Whether to cycle beyond num_steps
            (default: false).

            Inverse time decay:

decay_rate (float) – Decay rate
            (required).

staircase (bool) – Whether to apply decay in a discrete
            staircase, as opposed to continuous, fashion.
            (default: false).

            Cosine decay:

alpha (float) – Minimum learning rate value as a fraction of
            learning_rate
            (default: 0.0).

            Cosine decay with restarts:

t_mul (float) – Used to derive the number of iterations in the
            i-th period
            (default: 2.0).

m_mul (float) – Used to derive the initial learning rate of the
            i-th period
            (default: 1.0).

alpha (float) – Minimum learning rate value as a fraction of
            the learning_rate
            (default: 0.0).

            Linear cosine decay:

num_periods (float) – Number of periods in the cosine part of
            the decay
            (default: 0.5).

alpha (float) – Alpha value
            (default: 0.0).

beta (float) – Beta value
            (default: 0.001).

            Noisy linear cosine decay:

initial_variance (float) – Initial variance for the noise
            (default: 1.0).

variance_decay (float) – Decay for the noise's variance
            (default: 0.55).

num_periods (float) – Number of periods in the cosine part of
            the decay
            (default: 0.5).

alpha (float) – Alpha value
            (default: 0.0).

beta (float) – Beta value
            (default: 0.001).

        name (string): internal use.
        dtype (type): internal use.
        min_value (dtype-compatible value): internal use.
        max_value (dtype-compatible value): internal use.
    """

    def __init__(
        self, *, decay, unit, num_steps, initial_value, increasing=False, inverse=False, scale=1.0,
        name=None, dtype=None, min_value=None, max_value=None, **kwargs
    ):
        assert decay in (
            'cosine', 'cosine_restarts', 'exponential', 'inverse_time', 'linear', 'linear_cosine',
            'linear_cosine_noisy', 'polynomial'
        )
        assert unit in ('timesteps', 'episodes', 'updates')
        assert (isinstance(num_steps, int) or num_steps % 10.0 == 0.0) and num_steps > 0
        assert isinstance(initial_value, (floa

        if isinstance(initial_value, int):
            if dtype != 'int':
                raise TensorforceError.dtype(
                    name='Decaying', argument='initial_value', dtype=type(initial_value)
                )
        elif isinstance(initial_value, float):
            if dtype != 'float':
                raise TensorforceError.dtype(
                    name='Decaying', argument='initial_value', dtype=type(initial_value)
                )
        else:
            raise TensorforceError.unexpected()

        if decay == 'linear':
            assert len(kwargs) == 1 and 'final_value' in kwargs
            decay = 'polynomial'

        self.decay = decay
        self.num_steps = int(num_steps)
        self.initial_value = initial_value
        self.increasing = increasing
        self.inverse = inverse
        self.scale = scale
        self.kwargs = kwargs

        super().__init__(
            unit=unit, name=name, dtype=dtype, min_value=min_value, max_value=max_value
        )

    def min_value(self):
        if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('alpha', 0.0)
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= 0.0:
                min_value = 0.0
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('beta', 0.001)
                max_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
            else:
                min_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
                max_value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= self.kwargs['final_value']:
                min_value = self.kwargs['final_value']
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.kwargs['final_value']

        assert min_value <= max_value

        if self.increasing:
            assert 0.0 <= min_value <= max_value <= 1.0
            min_value, max_value = 1.0 - max_value, 1.0 - min_value

        if self.inverse:
            assert util.epsilon <= min_value <= max_value
            min_value, max_value = 1.0 / max_value, 1.0 / min_value

        if self.scale == 1.0:
            pass
        elif self.scale >= 0.0:
            min_value, max_value = self.scale * min_value, self.scale * max_value
        else:
            min_value, max_value = self.scale * max_value, self.scale * min_value

        return self.spec.py_type()(min_value)

    def max_value(self):
        if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('alpha', 0.0)
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= 0.0:
                min_value = 0.0
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            if self.initial_value >= 0.0:
                min_value = self.initial_value * self.kwargs.get('beta', 0.001)
                max_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
            else:
                min_value = self.initial_value * (
                    1.0 + self.kwargs.get('alpha', 0.0) + self.kwargs.get('beta', 0.001)
                )
                max_value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                min_value = max_value = self.initial_value
            elif self.initial_value >= self.kwargs['final_value']:
                min_value = self.kwargs['final_value']
                max_value = self.initial_value
            else:
                min_value = self.initial_value
                max_value = self.kwargs['final_value']

        assert min_value <= max_value

        if self.increasing:
            assert 0.0 <= min_value <= max_value <= 1.0
            min_value, max_value = 1.0 - max_value, 1.0 - min_value

        if self.inverse:
            assert 0.0 < min_value <= max_value
            min_value, max_value = 1.0 / max_value, 1.0 / min_value

        if self.scale == 1.0:
            pass
        elif self.scale >= 0.0:
            min_value, max_value = self.scale * min_value, self.scale * max_value
        else:
            min_value, max_value = self.scale * max_value, self.scale * min_value

        return self.spec.py_type()(max_value)

    def final_value(self):
        if self.decay == 'cosine' or self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            value = self.initial_value * self.kwargs.get('alpha', 0.0)

        elif self.decay == 'exponential' or self.decay == 'inverse_time':
            assert 0.0 <= self.kwargs['decay_rate'] <= 1.0
            if self.kwargs['decay_rate'] == 1.0:
                value = self.initial_value
            else:
                value = 0.0

        elif self.decay == 'linear_cosine' or self.decay == 'linear_cosine_noisy':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            assert 0.0 <= self.kwargs.get('beta', 0.0) <= 1.0
            value = self.initial_value * self.kwargs.get('beta', 0.001)

        elif self.decay == 'polynomial':
            if self.kwargs.get('power', 1.0) == 0.0:
                value = self.initial_value
            else:
                value = self.kwargs['final_value']

        if self.increasing:
            assert 0.0 <= value <= 1.0
            value = 1.0 - value

        if self.inverse:
            assert value > 0.0
            value = 1.0 / value

        if self.scale != 1.0:
            value = value * self.scale

        return self.spec.py_type()(value)

    def parameter_value(self, *, step):
        initial_value = tf_util.constant(value=self.initial_value, dtype='float')

        if self.decay == 'cosine':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            parameter = tf.keras.experimental.CosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                alpha=self.kwargs.get('alpha', 0.0)
            )(step=step)

        elif self.decay == 'cosine_restarts':
            assert 0.0 <= self.kwargs.get('alpha', 0.0) <= 1.0
            parameter = tf.keras.experimental.CosineDecayRestarts(
                initial_learning_rate=initial_value, first_decay_steps=(self.num_steps + 1),
                t_mul=self.kwargs.get('t_mul', 2.0), m_mul=self.kwargs.get('m_mul', 1.0),
                alpha=self.kwargs.get('alpha', 0.0)
            )(step=step)

        elif self.decay == 'exponential':
            assert self.kwargs['decay_rate'] >= 0.0
            parameter = tf.keras.optimizers.schedules.ExponentialDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                decay_rate=self.kwargs['decay_rate'], staircase=self.kwargs.get('staircase', False)
            )(step=step)

        elif self.decay == 'inverse_time':
            assert self.kwargs['decay_rate'] >= 0.0
            parameter = tf.keras.optimizers.schedules.InverseTimeDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                decay_rate=self.kwargs['decay_rate'], staircase=self.kwargs.get('staircase', False)
            )(step=step)

        elif self.decay == 'linear_cosine':
            assert self.kwargs.get('beta', 0.001) >= 0.0
            parameter = tf.keras.experimental.LinearCosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                num_periods=self.kwargs.get('num_periods', 0.5),
                alpha=self.kwargs.get('alpha', 0.0), beta=self.kwargs.get('beta', 0.001)
            )(step=step)

        elif self.decay == 'linear_cosine_noisy':
            assert self.kwargs.get('beta', 0.001) >= 0.0
            parameter = tf.keras.experimental.NoisyLinearCosineDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                initial_variance=self.kwargs.get('initial_variance', 1.0),
                variance_decay=self.kwargs.get('variance_decay', 0.55),
                num_periods=self.kwargs.get('num_periods', 0.5),
                alpha=self.kwargs.get('alpha', 0.0), beta=self.kwargs.get('beta', 0.001)
            )(step=step)

        elif self.decay == 'polynomial':
            assert self.kwargs.get('power', 1.0) >= 0.0
            parameter = tf.keras.optimizers.schedules.PolynomialDecay(
                initial_learning_rate=initial_value, decay_steps=(self.num_steps + 1),
                end_learning_rate=self.kwargs['final_value'], power=self.kwargs.get('power', 1.0),
                cycle=self.kwargs.get('cycle', False)
            )(step=step)

        if self.increasing:
            one = tf_util.constant(value=1.0, dtype='float')
            assertions = list()
            if self.config.create_tf_assertions:
                zero = tf_util.constant(value=0.0, dtype='float')
                assertions.append(tf.debugging.assert_greater_equal(x=parameter, y=zero))
                assertions.append(tf.debugging.assert_less_equal(x=parameter, y=one))
            with tf.control_dependencies(control_inputs=assertions):
                parameter = one - parameter

        if self.inverse:
            zero = tf_util.constant(value=0.0, dtype='float')
            epsilon = tf_util.constant(value=util.epsilon, dtype='float')
            parameter = tf.where(
                condition=(parameter > zero), x=(parameter + epsilon), y=(parameter - epsilon)
            )
            parameter = tf.math.reciprocal(x=parameter)

        if self.scale != 1.0:
            scale = tf_util.constant(value=self.scale, dtype='float')
            parameter = parameter * scale

        parameter = tf_util.cast(x=parameter, dtype=self.spec.type)

        return parameter


def compute_accuracy(self, dataloader: DataLoader):

import os
import random
from typing import List, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
from torch import optim, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from tqdm import tqdm
from firelab.config import Config

from src.utils.data_utils import construct_output_mask, flatten, remap_targets
from src.dataloaders.utils import create_custom_dataset
from src.utils.training_utils import (
    construct_optimizer,
    prune_logits,
    construct_per_group_optimizer,
    decrease_lr_in_optim_config
)

class TaskTrainer:
    def __init__(self, main_trainer: "LLLTrainer", task_idx:int):
        self.task_idx = task_idx
        self.main_trainer = main_trainer
        self.device_name = main_trainer.device_name
        self.config = main_trainer.config
        self.start_task_idx = self.config.get('start_task', 0)

        self.init_models()
        self.init_writer()

        self.criterion = nn.CrossEntropyLoss()
        self.optim = self.construct_optimizer()
        self.attrs = self.model.attrs if hasattr(self.model, 'attrs') else None
        self.task_ds_train, self.task_ds_test = main_trainer.data_splits[task_idx]
        self.output_mask = construct_output_mask(main_trainer.class_splits[task_idx], self.config.lll_setup.num_classes)
        self.classes = self.main_trainer.class_splits[self.task_idx]
        self.learned_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx])).tolist()
        self.learned_classes_mask = construct_output_mask(self.learned_classes, self.config.data.num_classes)
        self.seen_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx + 1])).tolist()
        self.seen_classes_mask = construct_output_mask(self.seen_classes, self.config.data.num_classes)
        if self.task_idx >= self.config.start_task:
            self.curr_classes_across_seen_mask = construct_output_mask(remap_targets(self.classes, self.seen_classes), len(self.seen_classes))
        self.init_dataloaders()
        self.init_episodic_memory()
        self.test_acc_batch_history = []
        self.after_iter_done_callbacks = []
        self.num_iters_done = 0
        self.num_epochs_done = 0

        self._after_init_hook()

    def init_models(self):
        self.model = self.main_trainer.model

        if self.config.hp.get('reinit_after_each_task'):
            self.model.load_state_dict(self.main_trainer.create_model().state_dict())

    def init_writer(self):
        if not self.config.get('no_saving'):
            self.writer = SummaryWriter(os.path.join(self.main_trainer.paths.logs_path, f'task_{self.task_idx}'), flush_secs=5)

    def clip_grad(self):
        if self.config.hp.get('clip_grad.value', -1) > 0:
            grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.config.hp.clip_grad.value)
            if not self.config.get('no_saving'):
                self.writer.add_scalar('cls/grad_norm', grad_norm, self.num_iters_done)

    def construct_optimizer(self):
        if self.config.hp.optim.get('reuse') and self.task_idx > 0:
            return self.get_previous_trainer().optim

        optim_conf = decrease_lr_in_optim_config(self.config.hp.optim, self.task_idx - self.config.get('start_task', 0))

        if optim_conf.has('groups'):
            return construct_per_group_optimizer(self.model, optim_conf)
        else:
            return construct_optimizer(self.model.parameters(), optim_conf)

    def init_dataloaders(self):
        self.train_dataloader = self.create_dataloader(self.task_ds_train, shuffle=True)
        self.test_dataloader = self.create_dataloader(self.task_ds_test, shuffle=False)

    def create_dataloader(self, dataset: List[Tuple[Any, int]], shuffle: bool, batch_size: int=None):
        if batch_size is None:
            batch_size = self.config.hp.batch_size

        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                          collate_fn=lambda b: list(zip(*b)), num_workers=4)

    def load_dataset(self, dataset: List[Tuple[Any, int]]):
        dl = self.create_dataloader(dataset, shuffle=False, batch_size=1)
        ds = [(x[0], y[0]) for x, y in dl]

        return ds

    def compute_loss(self, model: nn.Module, batch: Tuple[Tensor, Tensor]):
        if self.config.hp.use_class_attrs:
            targets = remap_targets(batch[1], self.seen_classes)
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(targets)).to(self.device_name)

            logits = model(x, attrs_mask=self.seen_classes_mask)

            if self.config.task_trainer == 'joint':
                pass
            else:
                logits = prune_logits(logits, self.curr_classes_across_seen_mask)
        else:
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(batch[1])).to(self.device_name)

            logits = model(x)
            logits = prune_logits(logits, self.output_mask)

        return self.criterion(logits, y)

    def _after_init_hook(self):
        pass

    def _before_train_hook(self):
        pass

    def _after_train_hook(self):
        pass

    def on_epoch_done(self):
        pass

    def run_after_iter_done_callbacks(self):
        for callback in self.after_iter_done_callbacks:
            callback(self)

    def get_previous_trainer(self) -> "TaskTrainer":
        if self.task_idx == 0 or (self.task_idx - 1) >= len(self.main_trainer.task_trainers):
            return None
        else:
            return self.main_trainer.task_trainers[self.task_idx - 1]

    def init_episodic_memory(self):
        if self.get_previous_trainer() is None:
            self.episodic_memory = []
        else:
            self.episodic_memory = self.get_previous_trainer().episodic_memory

    def update_episodic_memory(self):
        pass

    @property
    def is_trainable(self) -> bool:
        return self.task_idx >= self.config.start_task

    def start(self):
        """Runs training"""
        self._before_train_hook()

        assert self.is_trainable, "We do not have enough conditions to train this Task Trainer"\
                                  "(for example, previous trainers was not finished or this trainer was already run)"

        if self.task_idx == 0:
            num_epochs = self.config.hp.get('base_task_max_num_epochs', self.config.hp.max_num_epochs)
        else:
            num_epochs = self.config.hp.max_num_epochs

        epochs = range(1, num_epochs + 1)
        if self._should_tqdm_epochs(): epochs = tqdm(epochs, desc=f'Task #{self.task_idx}')

        for epoch in epochs:
            batches = self.train_dataloader

            if not self._should_tqdm_epochs():
                batches = tqdm(batches, desc=f'Task #{self.task_idx} [epoch {epoch}/{num_epochs}]')

            for batch in batches:
                self.train_on_batch(batch)
                self.num_iters_done += 1
                self.run_after_iter_done_callbacks()

                if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                    break

            self.num_epochs_done += 1
            self.on_epoch_done()

            if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                break

        self.update_episodic_memory()
        self._after_train_hook()

    def train_on_batch(self, batch):
        raise NotImplementedError

guessed = []
        self.model.eval()

        with torch.no_grad():
            for x, y in dataloader:
                x = torch.from_numpy(np.array(x)).to(self.device_name)
                y = torch.from_numpy(np.array(y)).to(self.device_name)

                pruned_logits = self.model.compute_pruned_predictions(x, self.output_mask)

                guessed.extend((pruned_logits.argmax(dim=1) == y).cpu().data.tolist())

        return np.mean(guessed)

    def compute_test_accuracy(self):
        return self.compute_accuracy(self.test_dataloader)

    def compute_train_accuracy(self):
        return self.compute_accuracy(self.train_dataloader)

    def _should_tqdm_epochs(self) -> bool:
        return self.config.hp.max_num_epochs > 10

    def sample_from_memory(self, batch_size: int) -> Tuple[Tensor, Tensor]:
        samples = random.choices(self.episodic_memory, k=batch_size)
        x, y = zip(*samples)
        x = torch.from_numpy(np.array(x)).to(self.device_name)
        y = torch.tensor(y).to(self.device_name)

        return x, y

    def sample_batch(self, dataset: List[Tuple[Any, int]], batch_size: int, replace: bool=False) -> Tuple[np.ndarray, np.ndarray]:
        batch_size = batch_size if replace else min(batch_size, len(dataset))
        idx = np.random.choice(range(len(dataset)), size=batch_size, replace=replace)
        x, y = zip(*[dataset[i] for i in idx])

        return x, y


import os
import random
from typing import List, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
from torch import optim, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from tqdm import tqdm
from firelab.config import Config

from src.utils.data_utils import construct_output_mask, flatten, remap_targets
from src.dataloaders.utils import create_custom_dataset
from src.utils.training_utils import (
    construct_optimizer,
    prune_logits,
    construct_per_group_optimizer,
    decrease_lr_in_optim_config
)

class TaskTrainer:
    def __init__(self, main_trainer: "LLLTrainer", task_idx:int):
        self.task_idx = task_idx
        self.main_trainer = main_trainer
        self.device_name = main_trainer.device_name
        self.config = main_trainer.config
        self.start_task_idx = self.config.get('start_task', 0)

        self.init_models()
        self.init_writer()

        self.criterion = nn.CrossEntropyLoss()
        self.optim = self.construct_optimizer()
        self.attrs = self.model.attrs if hasattr(self.model, 'attrs') else None
        self.task_ds_train, self.task_ds_test = main_trainer.data_splits[task_idx]
        self.output_mask = construct_output_mask(main_trainer.class_splits[task_idx], self.config.lll_setup.num_classes)
        self.classes = self.main_trainer.class_splits[self.task_idx]
        self.learned_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx])).tolist()
        self.learned_classes_mask = construct_output_mask(self.learned_classes, self.config.data.num_classes)
        self.seen_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx + 1])).tolist()
        self.seen_classes_mask = construct_output_mask(self.seen_classes, self.config.data.num_classes)
        if self.task_idx >= self.config.start_task:
            self.curr_classes_across_seen_mask = construct_output_mask(remap_targets(self.classes, self.seen_classes), len(self.seen_classes))
        self.init_dataloaders()
        self.init_episodic_memory()
        self.test_acc_batch_history = []
        self.after_iter_done_callbacks = []
        self.num_iters_done = 0
        self.num_epochs_done = 0

        self._after_init_hook()

    def init_models(self):
        self.model = self.main_trainer.model

        if self.config.hp.get('reinit_after_each_task'):
            self.model.load_state_dict(self.main_trainer.create_model().state_dict())

    def init_writer(self):
        if not self.config.get('no_saving'):
            self.writer = SummaryWriter(os.path.join(self.main_trainer.paths.logs_path, f'task_{self.task_idx}'), flush_secs=5)

    def clip_grad(self):
        if self.config.hp.get('clip_grad.value', -1) > 0:
            grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.config.hp.clip_grad.value)
            if not self.config.get('no_saving'):
                self.writer.add_scalar('cls/grad_norm', grad_norm, self.num_iters_done)

    def construct_optimizer(self):
        if self.config.hp.optim.get('reuse') and self.task_idx > 0:
            return self.get_previous_trainer().optim

        optim_conf = decrease_lr_in_optim_config(self.config.hp.optim, self.task_idx - self.config.get('start_task', 0))

        if optim_conf.has('groups'):
            return construct_per_group_optimizer(self.model, optim_conf)
        else:
            return construct_optimizer(self.model.parameters(), optim_conf)

    def init_dataloaders(self):
        self.train_dataloader = self.create_dataloader(self.task_ds_train, shuffle=True)
        self.test_dataloader = self.create_dataloader(self.task_ds_test, shuffle=False)

    def create_dataloader(self, dataset: List[Tuple[Any, int]], shuffle: bool, batch_size: int=None):
        if batch_size is None:
            batch_size = self.config.hp.batch_size

        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                          collate_fn=lambda b: list(zip(*b)), num_workers=4)

    def load_dataset(self, dataset: List[Tuple[Any, int]]):
        dl = self.create_dataloader(dataset, shuffle=False, batch_size=1)
        ds = [(x[0], y[0]) for x, y in dl]

        return ds

    def compute_loss(self, model: nn.Module, batch: Tuple[Tensor, Tensor]):
        if self.config.hp.use_class_attrs:
            targets = remap_targets(batch[1], self.seen_classes)
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(targets)).to(self.device_name)

            logits = model(x, attrs_mask=self.seen_classes_mask)

            if self.config.task_trainer == 'joint':
                pass
            else:
                logits = prune_logits(logits, self.curr_classes_across_seen_mask)
        else:
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(batch[1])).to(self.device_name)

            logits = model(x)
            logits = prune_logits(logits, self.output_mask)

        return self.criterion(logits, y)

    def _after_init_hook(self):
        pass

    def _before_train_hook(self):

    def _after_train_hook(self):
        pass

    def on_epoch_done(self):
        pass

    def run_after_iter_done_callbacks(self):
        for callback in self.after_iter_done_callbacks:
            callback(self)

    def get_previous_trainer(self) -> "TaskTrainer":
        if self.task_idx == 0 or (self.task_idx - 1) >= len(self.main_trainer.task_trainers):
            return None
        else:
            return self.main_trainer.task_trainers[self.task_idx - 1]

    def init_episodic_memory(self):
        if self.get_previous_trainer() is None:
            self.episodic_memory = []
        else:
            self.episodic_memory = self.get_previous_trainer().episodic_memory

    def update_episodic_memory(self):
        pass

    @property
    def is_trainable(self) -> bool:
        return self.task_idx >= self.config.start_task

    def start(self):
        """Runs training"""
        self._before_train_hook()

        assert self.is_trainable, "We do not have enough conditions to train this Task Trainer"\
                                  "(for example, previous trainers was not finished or this trainer was already run)"

        if self.task_idx == 0:
            num_epochs = self.config.hp.get('base_task_max_num_epochs', self.config.hp.max_num_epochs)
        else:
            num_epochs = self.config.hp.max_num_epochs

        epochs = range(1, num_epochs + 1)
        if self._should_tqdm_epochs(): epochs = tqdm(epochs, desc=f'Task #{self.task_idx}')

        for epoch in epochs:
            batches = self.train_dataloader

            if not self._should_tqdm_epochs():
                batches = tqdm(batches, desc=f'Task #{self.task_idx} [epoch {epoch}/{num_epochs}]')

            for batch in batches:
                self.train_on_batch(batch)
                self.num_iters_done += 1
                self.run_after_iter_done_callbacks()

                if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                    break

            self.num_epochs_done += 1
            self.on_epoch_done()

            if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                break

        self.update_episodic_memory()
        self._after_train_hook()

    def train_on_batch(self, batch):
        raise NotImplementedError

    def compute_accuracy(self, dataloader: DataLoader):
        guessed = []
        self.model.eval()

        with torch.no_grad():
            for x, y in dataloader:
                x = torch.from_numpy(np.array(x)).to(self.device_name)
                y = torch.from_numpy(np.array(y)).to(self.device_name)

                pruned_logits = self.model.compute_pruned_predictions(x, self.output_mask)

                guessed.extend((pruned_logits.argmax(dim=1) == y).cpu().data.tolist())

        return np.mean(guessed)

    def compute_test_accuracy(self):
        return self.compute_accuracy(self.test_dataloader)

    def compute_train_accuracy(self):
        return self.compute_accuracy(self.train_dataloader)

    def _should_tqdm_epochs(self) -> bool:
        return self.config.hp.max_num_epochs > 10

    def sample_from_memory(self, batch_size: int) -> Tuple[Tensor, Tensor]:
        samples = random.choices(self.episodic_memory, k=batch_size)
        x, y = zip(*samples)
        x = torch.from_numpy(np.array(x)).to(self.device_name)
        y = torch.tensor(y).to(self.device_name)

        return x, y

    def sample_batch(self, dataset: List[Tuple[Any, int]], batch_size: int, replace: bool=False) -> Tuple[np.ndarray, np.ndarray]:
        batch_size = batch_size if replace else min(batch_size, len(dataset))
        idx = np.random.choice(range(len(dataset)), size=batch_size, replace=replace)
        x, y = zip(*[dataset[i] for i in idx])

        return x, y


import os
import random
from typing import List, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
from torch import optim, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from tqdm import tqdm
from firelab.config import Config

from src.utils.data_utils import construct_output_mask, flatten, remap_targets
from src.dataloaders.utils import create_custom_dataset
from src.utils.training_utils import (
    construct_optimizer,
    prune_logits,
    construct_per_group_optimizer,
    decrease_lr_in_optim_config
)

class TaskTrainer:
    def __init__(self, main_trainer: "LLLTrainer", task_idx:int):
        self.task_idx = task_idx
        self.main_trainer = main_trainer
        self.device_name = main_trainer.device_name
        self.config = main_trainer.config
        self.start_task_idx = self.config.get('start_task', 0)

        self.init_models()
        self.init_writer()

        self.criterion = nn.CrossEntropyLoss()
        self.optim = self.construct_optimizer()
        self.attrs = self.model.attrs if hasattr(self.model, 'attrs') else None
        self.task_ds_train, self.task_ds_test = main_trainer.data_splits[task_idx]
        self.output_mask = construct_output_mask(main_trainer.class_splits[task_idx], self.config.lll_setup.num_classes)
        self.classes = self.main_trainer.class_splits[self.task_idx]
        self.learned_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx])).tolist()
        self.learned_classes_mask = construct_output_mask(self.learned_classes, self.config.data.num_classes)
        self.seen_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx + 1])).tolist()
        self.seen_classes_mask = construct_output_mask(self.seen_classes, self.config.data.num_classes)
        if self.task_idx >= self.config.start_task:
            self.curr_classes_across_seen_mask = construct_output_mask(remap_targets(self.classes, self.seen_classes), len(self.seen_classes))
        self.init_dataloaders()
        self.init_episodic_memory()
        self.test_acc_batch_history = []
        self.after_iter_done_callbacks = []
        self.num_iters_done = 0
        self.num_epochs_done = 0

        self._after_init_hook()

    def init_models(self):
        self.model = self.main_trainer.model

        if self.config.hp.get('reinit_after_each_task'):
            self.model.load_state_dict(self.main_trainer.create_model().state_dict())

    def init_writer(self):
        if not self.config.get('no_saving'):
            self.writer = SummaryWriter(os.path.join(self.main_trainer.paths.logs_path, f'task_{self.task_idx}'), flush_secs=5)

    def clip_grad(self):
        if self.config.hp.get('clip_grad.value', -1) > 0:
            grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.config.hp.clip_grad.value)
            if not self.config.get('no_saving'):
                self.writer.add_scalar('cls/grad_norm', grad_norm, self.num_iters_done)

    def construct_optimizer(self):
        if self.config.hp.optim.get('reuse') and self.task_idx > 0:
            return self.get_previous_trainer().optim

        optim_conf = decrease_lr_in_optim_config(self.config.hp.optim, self.task_idx - self.config.get('start_task', 0))

        if optim_conf.has('groups'):
            return construct_per_group_optimizer(self.model, optim_conf)
        else:
            return construct_optimizer(self.model.parameters(), optim_conf)

    def init_dataloaders(self):
        self.train_dataloader = self.create_dataloader(self.task_ds_train, shuffle=True)
        self.test_dataloader = self.create_dataloader(self.task_ds_test, shuffle=False)

    def create_dataloader(self, dataset: List[Tuple[Any, int]], shuffle: bool, batch_size: int=None):
        if batch_size is None:
            batch_size = self.config.hp.batch_size

        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                          collate_fn=lambda b: list(zip(*b)), num_workers=4)

    def load_dataset(self, dataset: List[Tuple[Any, int]]):
        dl = self.create_dataloader(dataset, shuffle=False, batch_size=1)
        ds = [(x[0], y[0]) for x, y in dl]

        return ds

    def compute_loss(self, model: nn.Module, batch: Tuple[Tensor, Tensor]):
        if self.config.hp.use_class_attrs:
            targets = remap_targets(batch[1], self.seen_classes)
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(targets)).to(self.device_name)

            logits = model(x, attrs_mask=self.seen_classes_mask)

            if self.config.task_trainer == 'joint':
                pass
            else:
                logits = prune_logits(logits, self.curr_classes_across_seen_mask)
        else:
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(batch[1])).to(self.device_name)

            logits = model(x)
            logits = prune_logits(logits, self.output_mask)

        return self.criterion(logits, y)

    def _after_init_hook(self):
        pass

    def _before_train_hook(self):
        pass

    def _after_train_hook(self):
        pass

    def on_epoch_done(self):
        pass

    def run_after_iter_done_callbacks(self):
        for callback in self.after_iter_done_callbacks:
            callback(self)

    def get_previous_trainer(self) -> "TaskTrainer":
        if self.task_idx == 0 or (self.task_idx - 1) >= len(self.main_trainer.task_trainers):
            return None
        else:
            return self.main_trainer.task_trainers[self.task_idx - 1]

    def init_episodic_memory(self):
        if self.get_previous_trainer() is None:
            self.episodic_memory = []
        else:
            self.episodic_memory = self.get_previous_trainer().episodic_memory

    def update_episodic_memory(self):
        pass

    @property
    def is_trainable(self) -> bool:
        return self.task_idx >= self.config.start_task

    def start(self):
        """Runs training"""
        self._before_train_hook()

        assert self.is_trainable, "We do not have enough conditions to train this Task Trainer"\
                                  "(for example, previous trainers was not finished or this trainer was already run)"

        if self.task_idx == 0:
            num_epochs = self.config.hp.get('base_task_max_num_epochs', self.config.hp.max_num_epochs)
        else:
            num_epochs = self.config.hp.max_num_epochs

        epochs = range(1, num_epochs + 1)
        if self._should_tqdm_epochs(): epochs = tqdm(epochs, desc=f'Task #{self.task_idx}')

        for epoch in epochs:
            batches = self.train_dataloader

            if not self._should_tqdm_epochs():
                batches = tqdm(batches, desc=f'Task #{self.task_idx} [epoch {epoch}/{num_epochs}]')

            for batch in batches:
                self.train_on_batch(batch)
                self.num_iters_done += 1
                self.run_after_iter_done_callbacks()

                if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                    break

            self.num_epochs_done += 1
            self.on_epoch_done()

            if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                break

        self.update_episodic_memory()
        self._after_train_hook()

    def train_on_batch(self, batch):
        raise NotImplementedError

    def compute_accuracy(self, dataloader: DataLoader):
        guessed = []
        self.model.eval()

        with torch.no_grad():
            for x, y in dataloader:
                x = torch.from_numpy(np.array(x)).to(self.device_name)
                y = torch.from_numpy(np.array(y)).to(self.device_name)

                pruned_logits = self.model.compute_pruned_predictions(x, self.output_mask)

                guessed.extend((pruned_logits.argmax(dim=1) == y).cpu().data.tolist())

        return np.mean(guessed)

    def compute_test_accuracy(self):
        return self.compute_accuracy(self.test_dataloader)

    def compute_train_accuracy(self):
        return self.compute_accuracy(self.train_dataloader)

    def _should_tqdm_epochs(self) -> bool:
        return self.config.hp.max_num_epochs > 10

    def sample_from_memory(self, batch_size: int) -> Tuple[Tensor, Tensor]:
        samples = random.choices(self.episodic_memory, k=batch_size)
        x, y = zip(*samples)
        x = torch.from_numpy(np.array(x)).to(self.device_name)
        y = torch.tensor(y).to(self.device_name)

        return x, y

    def sample_batch(self, dataset: List[Tuple[Any, int]], batch_size: int, replace: bool=False) -> Tuple[np.ndarray, np.ndarray]:
        batch_size = batch_size if replace

        idx = np.random.choice(range(len(dataset)), size=batch_size, replace=replace)
        x, y = zip(*[dataset[i] for i in idx])

        return x, y


for x, y in dataloader:
                x = torch.from_numpy(np.array(x)).to(self.device_name)
                y = torch.from_numpy(np.array(y)).to(self.device_name)

                pruned_logits = self.model.compute_pruned_predictions(x, self.output_mask)

                guessed.extend((pruned_logits.argmax(dim=1) == y).cpu().data.tolist())

import os
import random
from typing import List, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
from torch import optim, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from tqdm import tqdm
from firelab.config import Config

from src.utils.data_utils import construct_output_mask, flatten, remap_targets
from src.dataloaders.utils import create_custom_dataset
from src.utils.training_utils import (
    construct_optimizer,
    prune_logits,
    construct_per_group_optimizer,
    decrease_lr_in_optim_config
)

class TaskTrainer:
    def __init__(self, main_trainer: "LLLTrainer", task_idx:int):
        self.task_idx = task_idx
        self.main_trainer = main_trainer
        self.device_name = main_trainer.device_name
        self.config = main_trainer.config
        self.start_task_idx = self.config.get('start_task', 0)

        self.init_models()
        self.init_writer()

        self.criterion = nn.CrossEntropyLoss()
        self.optim = self.construct_optimizer()
        self.attrs = self.model.attrs if hasattr(self.model, 'attrs') else None
        self.task_ds_train, self.task_ds_test = main_trainer.data_splits[task_idx]
        self.output_mask = construct_output_mask(main_trainer.class_splits[task_idx], self.config.lll_setup.num_classes)
        self.classes = self.main_trainer.class_splits[self.task_idx]
        self.learned_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx])).tolist()
        self.learned_classes_mask = construct_output_mask(self.learned_classes, self.config.data.num_classes)
        self.seen_classes = np.unique(flatten(self.main_trainer.class_splits[self.start_task_idx:self.task_idx + 1])).tolist()
        self.seen_classes_mask = construct_output_mask(self.seen_classes, self.config.data.num_classes)
        if self.task_idx >= self.config.start_task:
            self.curr_classes_across_seen_mask = construct_output_mask(remap_targets(self.classes, self.seen_classes), len(self.seen_classes))
        self.init_dataloaders()
        self.init_episodic_memory()
        self.test_acc_batch_history = []
        self.after_iter_done_callbacks = []
        self.num_iters_done = 0
        self.num_epochs_done = 0

        self._after_init_hook()

    def init_models(self):
        self.model = self.main_trainer.model

        if self.config.hp.get('reinit_after_each_task'):
            self.model.load_state_dict(self.main_trainer.create_model().state_dict())

    def init_writer(self):
        if not self.config.get('no_saving'):
            self.writer = SummaryWriter(os.path.join(self.main_trainer.paths.logs_path, f'task_{self.task_idx}'), flush_secs=5)

    def clip_grad(self):
        if self.config.hp.get('clip_grad.value', -1) > 0:
            grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.config.hp.clip_grad.value)
            if not self.config.get('no_saving'):
                self.writer.add_scalar('cls/grad_norm', grad_norm, self.num_iters_done)

    def construct_optimizer(self):
        if self.config.hp.optim.get('reuse') and self.task_idx > 0:
            return self.get_previous_trainer().optim

        optim_conf = decrease_lr_in_optim_config(self.config.hp.optim, self.task_idx - self.config.get('start_task', 0))

        if optim_conf.has('groups'):
            return construct_per_group_optimizer(self.model, optim_conf)
        else:
            return construct_optimizer(self.model.parameters(), optim_conf)

    def init_dataloaders(self):
        self.train_dataloader = self.create_dataloader(self.task_ds_train, shuffle=True)
        self.test_dataloader = self.create_dataloader(self.task_ds_test, shuffle=False)

    def create_dataloader(self, dataset: List[Tuple[Any, int]], shuffle: bool, batch_size: int=None):
        if batch_size is None:
            batch_size = self.config.hp.batch_size

        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                          collate_fn=lambda b: list(zip(*b)), num_workers=4)

    def load_dataset(self, dataset: List[Tuple[Any, int]]):
        dl = self.create_dataloader(dataset, shuffle=False, batch_size=1)
        ds = [(x[0], y[0]) for x, y in dl]

        return ds

    def compute_loss(self, model: nn.Module, batch: Tuple[Tensor, Tensor]):
        if self.config.hp.use_class_attrs:
            targets = remap_targets(batch[1], self.seen_classes)
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(targets)).to(self.device_name)

            logits = model(x, attrs_mask=self.seen_classes_mask)

            if self.config.task_trainer == 'joint':
                pass
            else:
                logits = prune_logits(logits, self.curr_classes_across_seen_mask)
        else:
            x = torch.from_numpy(np.array(batch[0])).to(self.device_name)
            y = torch.from_numpy(np.array(batch[1])).to(self.device_name)

            logits = model(x)
            logits = prune_logits(logits, self.output_mask)

        return self.criterion(logits, y)

    def _after_init_hook(self):
        pass

    def _before_train_hook(self):
        pass

    def _after_train_hook(self):
        pass

    def on_epoch_done(self):
        pass

    def run_after_iter_done_callbacks(self):
        for callback in self.after_iter_done_callbacks:
            callback(self)

    def get_previous_trainer(self) -> "TaskTrainer":
        if self.task_idx == 0 or (self.task_idx - 1) >= len(self.main_trainer.task_trainers):
            return None
        else:
            return self.main_trainer.task_trainers[self.task_idx - 1]

    def init_episodic_memory(self):
        if self.get_previous_trainer() is None:
            self.episodic_memory = []
        else:
            self.episodic_memory = self.get_previous_trainer().episodic_memory

    def update_episodic_memory(self):
        pass

    @property
    def is_trainable(self) -> bool:
        return self.task_idx >= self.config.start_task

    def start(self):
        """Runs training"""
        self._before_train_hook()

        assert self.is_trainable, "We do not have enough conditions to train this Task Trainer"\
                                  "(for example, previous trainers was not finished or this trainer was already run)"

        if self.task_idx == 0:
            num_epochs = self.config.hp.get('base_task_max_num_epochs', self.config.hp.max_num_epochs)
        else:
            num_epochs = self.config.hp.max_num_epochs

        epochs = range(1, num_epochs + 1)
        if self._should_tqdm_epochs(): epochs = tqdm(epochs, desc=f'Task #{self.task_idx}')

        for epoch in epochs:
            batches = self.train_dataloader

            if not self._should_tqdm_epochs():
                batches = tqdm(batches, desc=f'Task #{self.task_idx} [epoch {epoch}/{num_epochs}]')

            for batch in batches:
                self.train_on_batch(batch)
                self.num_iters_done += 1
                self.run_after_iter_done_callbacks()

                if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                    break

            self.num_epochs_done += 1
            self.on_epoch_done()

            if self.num_iters_done >= self.config.hp.get('max_num_iters', 100000000):
                break

        self.update_episodic_memory()
        self._after_train_hook()

    def train_on_batch(self, batch):
        raise NotImplementedError

    def compute_accuracy(self, dataloader: DataLoader):
        guessed = []
        self.model.eval()

        with torch.no_grad():

        return np.mean(guessed)

    def compute_test_accuracy(self):
        return self.compute_accuracy(self.test_dataloader)

    def compute_train_accuracy(self):
        return self.compute_accuracy(self.train_dataloader)

    def _should_tqdm_epochs(self) -> bool:
        return self.config.hp.max_num_epochs > 10

    def sample_from_memory(self, batch_size: int) -> Tuple[Tensor, Tensor]:
        samples = random.choices(self.episodic_memory, k=batch_size)
        x, y = zip(*samples)
        x = torch.from_numpy(np.array(x)).to(self.device_name)
        y = torch.tensor(y).to(self.device_name)

        return x, y

    def sample_batch(self, dataset: List[Tuple[Any, int]], batch_size: int, replace: bool=False) -> Tuple[np.ndarray, np.ndarray]:
        batch_size = batch_size if replace else min(batch_size, len(dataset))
        idx = np.random.choice(range(len(dataset)), size=batch_size, replace=replace)
        x, y = zip(*[dataset[i] for i in idx])

        return x, y


# Appcelerator Titanium Module Packager
import os, subprocess, sys, glob, string, optparse, subprocess
import zipfile
from datetime import date

cwd = os.path.abspath(os.path.dirname(sys._getframe(0).f_code.co_filename))
os.chdir(cwd)
required_module_keys = ['name','version','moduleid','description','copyright','license','copyright','platform','minsdk']
module_defaults = {
	'description':'My module',
	'author': 'Your Name',
	'license' : 'Specify your license',
	'copyright' : 'Copyright (c) %s by Your Company' % str(date.today().year),
}
module_license_default = "TODO: place your license here and we'll include it in the module distribution"

def find_sdk(config):
	sdk = config['TITANIUM_SDK']
	return os.path.expandvars(os.path.expanduser(sdk))

def replace_vars(config,token):
	idx = token.find('$(')
	while idx != -1:
		idx2 = token.find(')',idx+2)
		if idx2 == -1: break
		key = token[idx+2:idx2]
		if not config.has_key(key): break
		token = token.replace('$(%s)' % key, config[key])
		idx = token.find('$(')
	return token

def read_ti_xcconfig():
	contents = open(os.path.join(cwd,'titanium.xcconfig')).read()
	config = {}
	for line in contents.splitlines(False):
		line = line.strip()
		if line[0:2]=='//': continue
		idx = line.find('=')
		if idx > 0:
			key = line[0:idx].strip()
			value = line[idx+1:].strip()
			config[key] = replace_vars(config,value)
	return config

def generate_doc(config):
	docdir = os.path.join(cwd,'documentation')
	if not os.path.exists(docdir):
		warn("Couldn't find documentation file at: %s" % docdir)
		return None

	try:
		import markdown2 as markdown
	except ImportError:
		import markdown
	documentation = []
	for file in os.listdir(docdir):
		if file in ignoreFiles or os.path.isdir(os.path.join(docdir, file)):
			continue
		md = open(os.path.join(docdir,file)).read()
		html = markdown.markdown(md)
		documentation.append({file:html});
	return documentation

def compile_js(manifest,config):
	js_file = os.path.join(cwd,'assets','__MODULE_ID__.js')
	if not os.path.exists(js_file): return

	from compiler import Compiler
	try:
		import json
	except:
		import simplejson as json

	compiler = Compiler(cwd, manifest['moduleid'], manifest['name'], 'commonjs')
	root_asset, module_assets = compiler.compile_module()

	root_asset_content = """
%s

	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[0]);
""" % root_asset

	module_asset_content = """
%s

	NSNumber *index = [map objectForKey:path];
	if (index == nil) {
		return nil;
	}
	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[index.integerValue]);
""" % module_assets

	from tools import splice_code

	assets_router = os.path.join(cwd,'Classes','___PROJECTNAMEASIDENTIFIER___ModuleAssets.m')
	splice_code(assets_router, 'asset', root_asset_content)
	splice_code(assets_router, 'resolve_asset', module_asset_content)

	# Generate the exports after crawling all of the available JS source
	exports = open('metadata.json','w')
	json.dump({'exports':compiler.exports }, exports)
	exports.close()

def die(msg):
	print msg
	sys.exit(1)

def info(msg):
	print "[INFO] %s" % msg

def warn(msg):
	print "[WARN] %s" % msg

def validate_license():
	c = open(os.path.join(cwd,'LICENSE')).read()
	if c.find(module_license_default)!=-1:
		warn('please update the LICENSE file with your license text before distributing')

def validate_manifest():
	path = os.path.join(cwd,'manifest')
	f = open(path)
	if not os.path.exists(path): die("missing %s" % path)
	manifest = {}
	for line in f.readlines():
		line = line.strip()
		if line[0:1]=='#': continue
		if line.find(':') < 0: continue
		key,value = line.split(':')
		manifest[key.strip()]=value.strip()
	for key in required_module_keys:
		if not manifest.has_key(key): die("missing required manifest key '%s'" % key)
		if module_defaults.has_key(key):
			defvalue = module_defaults[key]
			curvalue = manifest[key]
			if curvalue==defvalue: warn("please update the manifest key: '%s' to a non-default value" % key)
	return manifest,path

ignoreFiles = ['.DS_Store','.gitignore','libTitanium.a','titanium.jar','README']
ignoreDirs = ['.DS_Store','.svn','.git','CVSROOT']

def zip_dir(zf,dir,basepath,ignoreExt=[]):
	if not os.path.exists(dir): return
	for root, dirs, files in os.walk(dir):
		for name in ignoreDirs:
			if name in dirs:
				dirs.remove(name)	# don't visit ignored directories
		for file in files:
			if file in ignoreFiles: continue
			e = os.path.splitext(file)
			if len(e) == 2 and e[1] in ignoreExt: continue
			from_ = os.path.join(root, file)
			to_ = from_.replace(dir, '%s/%s'%(basepath,dir), 1)
			zf.write(from_, to_)

def glob_libfiles():
	files = []
	for libfile in glob.glob('build/**/*.a'):
		if libfile.find('Release-')!=-1:
			files.append(libfile)
	return files

def build_module(manifest,config):
	from tools import ensure_dev_path
	ensure_dev_path()

	rc = os.system("xcodebuild -sdk iphoneos -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
	rc = os.system("xcodebuild -sdk iphonesimulator -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
    # build the merged library using lipo
	moduleid = manifest['moduleid']
	libpaths = ''
	for libfile in glob_libfiles():
		libpaths+='%s ' % libfile

	os.system("lipo %s -create -output build/lib%s.a" %(libpaths,moduleid))

global options

	if options.skip_docs:
		info("Skipping documentation generation.")
		return False
	else:
		info("Module apidoc generation can be skipped using --skip-docs")
	apidoc_path = os.path.join(cwd, "apidoc")
	if not os.path.exists(apidoc_path):
		warn("Skipping apidoc generation. No apidoc folder found at: %s" % apidoc_path)
		return False

	if not os.path.exists(apidoc_build_path):
	    os.makedirs(apidoc_build_path)
	ti_root = string.strip(subprocess.check_output(["echo $TI_ROOT"], shell=True))
	if not len(ti_root) > 0:
		warn("Not generating documentation from the apidoc folder. The titanium_mobile repo could not be found.")
		warn("Set the TI_ROOT environment variable to the parent folder where the titanium_mobile repo resides (eg.'export TI_ROOT=/Path').")
		return False
	docgen = os.path.join(ti_root, "titanium_mobile", "apidoc", "docgen.py")
	if not os.path.exists(docgen):
		warn("Not generating documentation from the apidoc folder. Couldn't find docgen.py at: %s" % docgen)
		return False

	info("Generating documentation from the apidoc folder.")
	rc = os.system("\"%s\" --format=jsca,modulehtml --css=styles.css -o \"%s\" -e \"%s\"" % (docgen, apidoc_build_path, apidoc_path))
	if rc != 0:
		die("docgen failed")
	return True

def package_module(manifest,mf,config):
	name = manifest['name'].lower()
	moduleid = manifest['moduleid'].lower()
	version = manifest['version']
	modulezip = '%s-iphone-%s.zip' % (moduleid,version)
	if os.path.exists(modulezip): os.remove(modulezip)
	zf = zipfile.ZipFile(modulezip, 'w', zipfile.ZIP_DEFLATED)
	modulepath = 'modules/iphone/%s/%s' % (moduleid,version)
	zf.write(mf,'%s/manifest' % modulepath)
	libname = 'lib%s.a' % moduleid
	zf.write('build/%s' % libname, '%s/%s' % (modulepath,libname))
	docs = generate_doc(config)
	if docs!=None:
		for doc in docs:
			for file, html in doc.iteritems():
				filename = string.replace(file,'.md','.html')
				zf.writestr('%s/documentation/%s'%(modulepath,filename),html)

	apidoc_build_path = os.path.join(cwd, "build", "apidoc")
	if generate_apidoc(apidoc_build_path):
		for file in os.listdir(apidoc_build_path):
			if file in ignoreFiles or os.path.isdir(os.path.join(apidoc_build_path, file)):
				continue
			zf.write(os.path.join(apidoc_build_path, file), '%s/documentation/apidoc/%s' % (modulepath, file))

	zip_dir(zf,'assets',modulepath,['.pyc','.js'])
	zip_dir(zf,'example',modulepath,['.pyc'])
	zip_dir(zf,'platform',modulepath,['.pyc','.js'])
	zf.write('LICENSE','%s/LICENSE' % modulepath)
	zf.write('module.xcconfig','%s/module.xcconfig' % modulepath)
	exports_file = 'metadata.json'
	if os.path.exists(exports_file):
		zf.write(exports_file, '%s/%s' % (modulepath, exports_file))
	zf.close()

if __name__ == '__main__':
	global options

	parser = optparse.OptionParser()
	parser.add_option("-s", "--skip-docs",
			dest="skip_docs",
			action="store_true",
			help="Will skip building documentation in apidoc folder",
			default=False)
	(options, args) = parser.parse_args()

	manifest,mf = validate_manifest()
	validate_license()
	config = read_ti_xcconfig()

	sdk = find_sdk(config)
	sys.path.insert(0,os.path.join(sdk,'iphone'))
	sys.path.append(os.path.join(sdk, "common"))

	compile_js(manifest,config)
	build_module(manifest,config)
	package_module(manifest,mf,config)
	sys.exit(0)


# Appcelerator Titanium Module Packager
import os, subprocess, sys, glob, string, optparse, subprocess
import zipfile
from datetime import date

cwd = os.path.abspath(os.path.dirname(sys._getframe(0).f_code.co_filename))
os.chdir(cwd)
required_module_keys = ['name','version','moduleid','description','copyright','license','copyright','platform','minsdk']
module_defaults = {
	'description':'My module',
	'author': 'Your Name',
	'license' : 'Specify your license',
	'copyright' : 'Copyright (c) %s by Your Company' % str(date.today().year),
}
module_license_default = "TODO: place your license here and we'll include it in the module distribution"

def find_sdk(config):
	sdk = config['TITANIUM_SDK']
	return os.path.expandvars(os.path.expanduser(sdk))

def replace_vars(config,token):
	idx = token.find('$(')
	while idx != -1:
		idx2 = token.find(')',idx+2)
		if idx2 == -1: break
		key = token[idx+2:idx2]
		if not config.has_key(key): break
		token = token.replace('$(%s)' % key, config[key])
		idx = token.find('$(')
	return token

def read_ti_xcconfig():
	contents = open(os.path.join(cwd,'titanium.xcconfig')).read()
	config = {}
	for line in contents.splitlines(False):
		line = line.strip()
		if line[0:2]=='//': continue
		idx = line.find('=')
		if idx > 0:
			key = line[0:idx].strip()
			value = line[idx+1:].strip()
			config[key] = replace_vars(config,value)
	return config

def generate_doc(config):
	docdir = os.path.join(cwd,'documentation')
	if not os.path.exists(docdir):
		warn("Couldn't find documentation file at: %s" % docdir)
		return None

	try:
		import markdown2 as markdown
	except ImportError:
		import markdown
	documentation = []
	for file in os.listdir(docdir):
		if file in ignoreFiles or os.path.isdir(os.path.join(docdir, file)):
			continue
		md = open(os.path.join(docdir,file)).read()
		html = markdown.markdown(md)
		documentation.append({file:html});
	return documentation

def compile_js(manifest,config):
	js_file = os.path.join(cwd,'assets','__MODULE_ID__.js')
	if not os.path.exists(js_file): return

	from compiler import Compiler
	try:
		import json
	except:
		import simplejson as json

	compiler = Compiler(cwd, manifest['moduleid'], manifest['name'], 'commonjs')
	root_asset, module_assets = compiler.compile_module()

	root_asset_content = """
%s

	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[0]);
""" % root_asset

	module_asset_content = """
%s

	NSNumber *index = [map objectForKey:path];
	if (index == nil) {
		return nil;
	}
	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[index.integerValue]);
""" % module_assets

	from tools import splice_code

	assets_router = os.path.join(cwd,'Classes','___PROJECTNAMEASIDENTIFIER___ModuleAssets.m')
	splice_code(assets_router, 'asset', root_asset_content)
	splice_code(assets_router, 'resolve_asset', module_asset_content)

	# Generate the exports after crawling all of the available JS source
	exports = open('metadata.json','w')
	json.dump({'exports':compiler.exports }, exports)
	exports.close()

def die(msg):

def info(msg):
	print "[INFO] %s" % msg

def warn(msg):
	print "[WARN] %s" % msg

def validate_license():
	c = open(os.path.join(cwd,'LICENSE')).read()
	if c.find(module_license_default)!=-1:
		warn('please update the LICENSE file with your license text before distributing')

def validate_manifest():
	path = os.path.join(cwd,'manifest')
	f = open(path)
	if not os.path.exists(path): die("missing %s" % path)
	manifest = {}
	for line in f.readlines():
		line = line.strip()
		if line[0:1]=='#': continue
		if line.find(':') < 0: continue
		key,value = line.split(':')
		manifest[key.strip()]=value.strip()
	for key in required_module_keys:
		if not manifest.has_key(key): die("missing required manifest key '%s'" % key)
		if module_defaults.has_key(key):
			defvalue = module_defaults[key]
			curvalue = manifest[key]
			if curvalue==defvalue: warn("please update the manifest key: '%s' to a non-default value" % key)
	return manifest,path

ignoreFiles = ['.DS_Store','.gitignore','libTitanium.a','titanium.jar','README']
ignoreDirs = ['.DS_Store','.svn','.git','CVSROOT']

def zip_dir(zf,dir,basepath,ignoreExt=[]):
	if not os.path.exists(dir): return
	for root, dirs, files in os.walk(dir):
		for name in ignoreDirs:
			if name in dirs:
				dirs.remove(name)	# don't visit ignored directories
		for file in files:
			if file in ignoreFiles: continue
			e = os.path.splitext(file)
			if len(e) == 2 and e[1] in ignoreExt: continue
			from_ = os.path.join(root, file)
			to_ = from_.replace(dir, '%s/%s'%(basepath,dir), 1)
			zf.write(from_, to_)

def glob_libfiles():
	files = []
	for libfile in glob.glob('build/**/*.a'):
		if libfile.find('Release-')!=-1:
			files.append(libfile)
	return files

def build_module(manifest,config):
	from tools import ensure_dev_path
	ensure_dev_path()

	rc = os.system("xcodebuild -sdk iphoneos -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
	rc = os.system("xcodebuild -sdk iphonesimulator -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
    # build the merged library using lipo
	moduleid = manifest['moduleid']
	libpaths = ''
	for libfile in glob_libfiles():
		libpaths+='%s ' % libfile

	os.system("lipo %s -create -output build/lib%s.a" %(libpaths,moduleid))

def generate_apidoc(apidoc_build_path):
	global options

	if options.skip_docs:
		info("Skipping documentation generation.")
		return False
	else:
		info("Module apidoc generation can be skipped using --skip-docs")
	apidoc_path = os.path.join(cwd, "apidoc")
	if not os.path.exists(apidoc_path):
		warn("Skipping apidoc generation. No apidoc folder found at: %s" % apidoc_path)
		return False

	if not os.path.exists(apidoc_build_path):
	    os.makedirs(apidoc_build_path)
	ti_root = string.strip(subprocess.check_output(["echo $TI_ROOT"], shell=True))
	if not len(ti_root) > 0:
		warn("Not generating documentation from the apidoc folder. The titanium_mobile repo could not be found.")
		warn("Set the TI_ROOT environment variable to the parent folder where the titanium_mobile repo resides (eg.'export TI_ROOT=/Path').")
		return False
	docgen = os.path.join(ti_root, "titanium_mobile", "apidoc", "docgen.py")
	if not os.path.exists(docgen):
		warn("Not generating documentation from the apidoc folder. Couldn't find docgen.py at: %s" % docgen)
		return False

	info("Generating documentation from the apidoc folder.")
	rc = os.system("\"%s\" --format=jsca,modulehtml --css=styles.css -o \"%s\" -e \"%s\"" % (docgen, apidoc_build_path, apidoc_path))
	if rc != 0:
		die("docgen failed")
	return True

def package_module(manifest,mf,config):
	name = manifest['name'].lower()
	moduleid = manifest['moduleid'].lower()
	version = manifest['version']
	modulezip = '%s-iphone-%s.zip' % (moduleid,version)
	if os.path.exists(modulezip): os.remove(modulezip)
	zf = zipfile.ZipFile(modulezip, 'w', zipfile.ZIP_DEFLATED)
	modulepath = 'modules/iphone/%s/%s' % (moduleid,version)
	zf.write(mf,'%s/manifest' % modulepath)
	libname = 'lib%s.a' % moduleid
	zf.write('build/%s' % libname, '%s/%s' % (modulepath,libname))
	docs = generate_doc(config)
	if docs!=None:
		for doc in docs:
			for file, html in doc.iteritems():
				filename = string.replace(file,'.md','.html')
				zf.writestr('%s/documentation/%s'%(modulepath,filename),html)

	apidoc_build_path = os.path.join(cwd, "build", "apidoc")
	if generate_apidoc(apidoc_build_path):
		for file in os.listdir(apidoc_build_path):
			if file in ignoreFiles or os.path.isdir(os.path.join(apidoc_build_path, file)):
				continue
			zf.write(os.path.join(apidoc_build_path, file), '%s/documentation/apidoc/%s' % (modulepath, file))

	zip_dir(zf,'assets',modulepath,['.pyc','.js'])
	zip_dir(zf,'example',modulepath,['.pyc'])
	zip_dir(zf,'platform',modulepath,['.pyc','.js'])
	zf.write('LICENSE','%s/LICENSE' % modulepath)
	zf.write('module.xcconfig','%s/module.xcconfig' % modulepath)
	exports_file = 'metadata.json'
	if os.path.exists(exports_file):
		zf.write(exports_file, '%s/%s' % (modulepath, exports_file))
	zf.close()

if __name__ == '__main__':
	global options

	parser = optparse.OptionParser()
	parser.add_option("-s", "--skip-docs",
			dest="skip_docs",
			action="store_true",
			help="Will skip building documentation in apidoc folder",
			default=False)
	(options, args) = parser.parse_args()

	manifest,mf = validate_manifest()
	validate_license()
	config = read_ti_xcconfig()

	sdk = find_sdk(config)
	sys.path.insert(0,os.path.join(sdk,'iphone'))
	sys.path.append(os.path.join(sdk, "common"))

	compile_js(manifest,config)
	build_module(manifest,config)
	package_module(manifest,mf,config)
	sys.exit(0)


# Appcelerator Titanium Module Packager
import os, subprocess, sys, glob, string, optparse, subprocess
import zipfile
from datetime import date

cwd = os.path.abspath(os.path.dirname(sys._getframe(0).f_code.co_filename))
os.chdir(cwd)
required_module_keys = ['name','version','moduleid','description','copyright','license','copyright','platform','minsdk']
module_defaults = {
	'description':'My module',
	'author': 'Your Name',
	'license' : 'Specify your license',
	'copyright' : 'Copyright (c) %s by Your Company' % str(date.today().year),
}
module_license_default = "TODO: place your license here and we'll include it in the module distribution"

def find_sdk(config):
	sdk = config['TITANIUM_SDK']
	return os.path.expandvars(os.path.expanduser(sdk))

def replace_vars(config,token):
	idx = token.find('$(')
	while idx != -1:
		idx2 = token.find(')',idx+2)
		if idx2 == -1: break
		key = token[idx+2:idx2]
		if not config.has_key(key): break
		token = token.replace('$(%s)' % key, config[key])
		idx = token.find('$(')
	return token

def read_ti_xcconfig():
	contents = open(os.path.join(cwd,'titanium.xcconfig')).read()
	config = {}
	for line in contents.splitlines(False):
		line = line.strip()
		if line[0:2]=='//': continue
		idx = line.find('=')
		if idx > 0:
			key = line[0:idx].strip()
			value = line[idx+1:].strip()
			config[key] = replace_vars(config,value)
	return config

def generate_doc(config):
	docdir = os.path.join(cwd,'documentation')
	if not os.path.exists(docdir):
		warn("Couldn't find documentation file at: %s" % docdir)
		return None

	try:
		import markdown2 as markdown
	except ImportError:
		import markdown
	documentation = []
	for file in os.listdir(docdir):
		if file in ignoreFiles or os.path.isdir(os.path.join(docdir, file)):
			continue
		md = open(os.path.join(docdir,file)).read()
		html = markdown.markdown(md)
		documentation.append({file:html});
	return documentation

def compile_js(manifest,config):
	js_file = os.path.join(cwd,'assets','__MODULE_ID__.js')
	if not os.path.exists(js_file): return

	from compiler import Compiler
	try:
		import json
	except:
		import simplejson as json

	compiler = Compiler(cwd, manifest['moduleid'], manifest['name'], 'commonjs')
	root_asset, module_assets = compiler.compile_module()

	root_asset_content = """
%s

	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[0]);
""" % root_asset

	module_asset_content = """
%s

	NSNumber *index = [map objectForKey:path];
	if (index == nil) {
		return nil;
	}
	return filterDataInRange([NSData dataWithBytesNoCopy:data length:sizeof(data) freeWhenDone:NO], ranges[index.integerValue]);
""" % module_assets

	from tools import splice_code

	assets_router = os.path.join(cwd,'Classes','___PROJECTNAMEASIDENTIFIER___ModuleAssets.m')
	splice_code(assets_router, 'asset', root_asset_content)
	splice_code(assets_router, 'resolve_asset', module_asset_content)

	# Generate the exports after crawling all of the available JS source
	exports = open('metadata.json','w')
	json.dump({'exports':compiler.exports }, exports)
	exports.close()

def die(msg):
	print msg
	sys.exit(1)

def info(msg):
	print "[INFO] %s" % msg

def warn(msg):
	print "[WARN] %s" % msg

def validate_license():
	c = open(os.path.join(cwd,'LICENSE')).read()
	if c.find(module_license_default)!=-1:
		warn('please update the LICENSE file with your license text before distributing')

def validate_manifest():
	path = os.path.join(cwd,'manifest')
	f = open(path)
	if not os.path.exists(path): die("missing %s" % path)
	manifest = {}
	for line in f.readlines():
		line = line.strip()
		if line[0:1]=='#': continue
		if line.find(':') < 0: continue
		key,value = line.split(':')
		manifest[key.strip()]=value.strip()
	for key in required_module_keys:
		if not manifest.has_key(key): die("missing required manifest key '%s'" % key)
		if module_defaults.has_key(key):
			defvalue = module_defaults[key]
			curvalue = manifest[key]
			if curvalue==defvalue: warn("please update the manifest key: '%s' to a non-default value" % key)
	return manifest,path

ignoreFiles = ['.DS_Store','.gitignore','libTitanium.a','titanium.jar','README']
ignoreDirs = ['.DS_Store','.svn','.git','CVSROOT']

def zip_dir(zf,dir,basepath,ignoreExt=[]):
	if not os.path.exists(dir): return
	for root, dirs, files in os.walk(dir):
		for name in ignoreDirs:
			if name in dirs:
				dirs.remove(name)	# don't visit ignored directories
		for file in files:
			if file in ignoreFiles: continue
			e = os.path.splitext(file)
			if len(e) == 2 and e[1] in ignoreExt: continue
			from_ = os.path.join(root, file)
			to_ = from_.replace(dir, '%s/%s'%(basepath,dir), 1)
			zf.write(from_, to_)

def glob_libfiles():
	files = []
	for libfile in glob.glob('build/**/*.a'):
		if libfile.find('Release-')!=-1:
			files.append(libfile)
	return files

def build_module(manifest,config):
	from too

	ensure_dev_path()

	rc = os.system("xcodebuild -sdk iphoneos -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
	rc = os.system("xcodebuild -sdk iphonesimulator -configuration Release")
	if rc != 0:
		die("xcodebuild failed")
    # build the merged library using lipo
	moduleid = manifest['moduleid']
	libpaths = ''
	for libfile in glob_libfiles():
		libpaths+='%s ' % libfile

	os.system("lipo %s -create -output build/lib%s.a" %(libpaths,moduleid))

def generate_apidoc(apidoc_build_path):
	global options

	if options.skip_docs:
		info("Skipping documentation generation.")
		return False
	else:
		info("Module apidoc generation can be skipped using --skip-docs")
	apidoc_path = os.path.join(cwd, "apidoc")
	if not os.path.exists(apidoc_path):
		warn("Skipping apidoc generation. No apidoc folder found at: %s" % apidoc_path)
		return False

	if not os.path.exists(apidoc_build_path):
	    os.makedirs(apidoc_build_path)
	ti_root = string.strip(subprocess.check_output(["echo $TI_ROOT"], shell=True))
	if not len(ti_root) > 0:
		warn("Not generating documentation from the apidoc folder. The titanium_mobile repo could not be found.")
		warn("Set the TI_ROOT environment variable to the parent folder where the titanium_mobile repo resides (eg.'export TI_ROOT=/Path').")
		return False
	docgen = os.path.join(ti_root, "titanium_mobile", "apidoc", "docgen.py")
	if not os.path.exists(docgen):
		warn("Not generating documentation from the apidoc folder. Couldn't find docgen.py at: %s" % docgen)
		return False

	info("Generating documentation from the apidoc folder.")
	rc = os.system("\"%s\" --format=jsca,modulehtml --css=styles.css -o \"%s\" -e \"%s\"" % (docgen, apidoc_build_path, apidoc_path))
	if rc != 0:
		die("docgen failed")
	return True

def package_module(manifest,mf,config):
	name = manifest['name'].lower()
	moduleid = manifest['moduleid'].lower()
	version = manifest['version']
	modulezip = '%s-iphone-%s.zip' % (moduleid,version)
	if os.path.exists(modulezip): os.remove(modulezip)
	zf = zipfile.ZipFile(modulezip, 'w', zipfile.ZIP_DEFLATED)
	modulepath = 'modules/iphone/%s/%s' % (moduleid,version)
	zf.write(mf,'%s/manifest' % modulepath)
	libname = 'lib%s.a' % moduleid
	zf.write('build/%s' % libname, '%s/%s' % (modulepath,libname))
	docs = generate_doc(config)
	if docs!=None:
		for doc in docs:
			for file, html in doc.iteritems():
				filename = string.replace(file,'.md','.html')
				zf.writestr('%s/documentation/%s'%(modulepath,filename),html)

	apidoc_build_path = os.path.join(cwd, "build", "apidoc")
	if generate_apidoc(apidoc_build_path):
		for file in os.listdir(apidoc_build_path):
			if file in ignoreFiles or os.path.isdir(os.path.join(apidoc_build_path, file)):
				continue
			zf.write(os.path.join(apidoc_build_path, file), '%s/documentation/apidoc/%s' % (modulepath, file))

	zip_dir(zf,'assets',modulepath,['.pyc','.js'])
	zip_dir(zf,'example',modulepath,['.pyc'])
	zip_dir(zf,'platform',modulepath,['.pyc','.js'])
	zf.write('LICENSE','%s/LICENSE' % modulepath)
	zf.write('module.xcconfig','%s/module.xcconfig' % modulepath)
	exports_file = 'metadata.json'
	if os.path.exists(exports_file):
		zf.write(exports_file, '%s/%s' % (modulepath, exports_file))
	zf.close()

if __name__ == '__main__':
	global options

	parser = optparse.OptionParser()
	parser.add_option("-s", "--skip-docs",
			dest="skip_docs",
			action="store_true",
			help="Will skip building documentation in apidoc folder",
			default=False)
	(options, args) = parser.parse_args()

	manifest,mf = validate_manifest()
	validate_license()
	config = read_ti_xcconfig()

	sdk = find_sdk(config)
	sys.path.insert(0,os.path.join(sdk,'iphone'))
	sys.path.append(os.path.join(sdk, "common"))

	compile_js(manifest,config)
	build_module(manifest,config)
	package_module(manifest,mf,config)
	sys.exit(0)