#
# Copyright 2018 Analytics Zoo Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import sys
import io
import torch
from importlib.util import find_spec
from zoo.pipeline.api.torch import zoo_pickle_module
from bigdl.optim.optimizer import OptimMethod

if find_spec('jep') is None:
    raise Exception("jep not found, please install jep first.")


class TorchOptim(OptimMethod):
    """
    TorchOptim wraps a torch optimizer for distributed inference or training.
    """

    def __init__(self, optim_bytes, bigdl_type="float"):
        """
        :param bigdl_type:
        """
        super(TorchOptim, self).__init__(None, bigdl_type, optim_bytes)

    @staticmethod
    def from_pytorch(optim):
        bys = io.BytesIO()
        torch.save(optim, bys, pickle_module=zoo_pickle_module)
        zoo_optim = TorchOptim(bys.getvalue())
        return zoo_optim

@@ -1010,16 +1010,27 @@
 m_bytes,
+ decayType,
  bigdl_t
@@ -1164,16 +1164,27 @@
 im_bytes
+, decayType
 )%0A%0A    @
@@ -1222,18 +1222,685 @@
 ch(optim
-):
+, decayType=%22EpochDecay%22):%0A        %22%22%22%0A        :param optim: Pytorch optimizer or LrScheduler.%0A        :param decayType: one string of EpochDecay, IterationDecay, EpochDecayByScore if%0A                          the optim is LrScheduler.%0A                          EpochDecay: call LrScheduler.step() every epoch.%0A                          IterationDecay: call LrScheduler.step() every iteration.%0A                          EpochDecayByScore: call LrScheduler.step(val_score) every epoch,%0A                              val_score is the return value of the first validation method.%0A                              ReduceLROnPlateau must use this EpochDecayByScore.%0A        %22%22%22
 %0A       
@@ -2028,16 +2028,27 @@
 tvalue()
+, decayType
 )%0A      

support epochdecay iterationdecay epochdecaybyscore in TorchOptim (#2951)