from ..compat import BytesIO
from ..neural._classes.model import Model

try:
    import torch.autograd
    import torch
except ImportError:
    pass


class PytorchWrapper(Model):
    '''Wrap a PyTorch model, so that it has the same API as Thinc models.
    To optimize the model, you'll need to create a PyTorch optimizer and call
    optimizer.step() after each batch --- see examples/wrap_pytorch.py
    '''
    def __init__(self, model):
        Model.__init__(self)
        self._model = model

    def begin_update(self, x_data, drop=0.):
        '''Return the output of the wrapped PyTorch model for the given input,
        along with a callback to handle the backward pass.
        '''
        x_var = torch.autograd.Variable(torch.Tensor(x_data),
                                        requires_grad=True)
        # Make prediction
        y_var = self._model(x_var)
        def backward_pytorch(dy_data, sgd=None):
            dy_var = torch.autograd.Variable(torch.Tensor(dy_data))
            torch.autograd.backward((y_var,), grad_variables=(dy_var,))
            dX = self.ops.asarray(x_var.grad.data)
            if sgd is not None:
                optimizer.step()
            return dX
        return self.ops.asarray(y_var.data), backward

    def to_disk(self, path):
        # TODO: Untested
        torch.save(self._model.state_dict(), str(path))

    def from_disk(self, path):
        # TODO: Untested
        self._model.load_state_dict(torch.load(path))

    def to_bytes(self):
        # TODO: Untested
        filelike = BytesIO()
        torch.save(self._model.state_dict(), filelike)
        return filelike.read()

    def from_bytes(self, data):
        # TODO: Untested
        filelike = BytesIO(data)
        self._model.load_state_dict(torch.load(filelike))

    def to_gpu(self, device_num):
        # TODO: Implement
        raise NotImplementedError

    def to_cpu(self):
        # TODO: Implement
        raise NotImplementedError

    def resize_output(self):
        # TODO: Required for spaCy add label
        raise NotImplementedError

    def resize_input(self):
        # TODO: Not required yet, but should be useful
        raise NotImplementedError

    @contextlib.contextmanager
    def use_params(self, params): # pragma: no cover
        # TODO: Implement
        raise NotImplementedError


@@ -1,12 +1,30 @@
+import contextlib%0A
 from ..compa
@@ -114,16 +114,39 @@
 utograd%0A
+    import torch.optim%0A
     impo
@@ -193,17 +193,17 @@
 class Py
-t
+T
 orchWrap
@@ -532,16 +532,47 @@
  = model
+%0A        self._optimizer = None
 %0A%0A    de
@@ -1235,100 +1235,658 @@
     
-optimizer.step()%0A            return dX%0A        return self.ops.asarray(y_var.data), backward
+if self._optimizer is None:%0A                    self._optimizer = self._create_optimizer(sgd)%0A                self._optimizer.step()%0A                self._optimizer.zero_grad()%0A            return dX%0A        return self.ops.asarray(y_var.data), backward_pytorch%0A%0A    def _create_optimizer(self, sgd):%0A        params = self._model.parameters()%0A        if sgd.b1 != 0 and sgd.b2 != 0:%0A            optimizer = torch.optim.Adam(params, lr=sgd.alpha, betas=(sgd.b1, sgd.b2))%0A        elif sgd.b2 == 0:%0A            optimizer = torch.optim.SGD(params, lr=sgd.alpha, momentum=sgd.b1)%0A        else:%0A            raise NotImplementedError%0A        return optimizer
 %0A%0A  
@@ -2465,142 +2465,85 @@
     
-# TODO: Implement%0A        raise NotImplementedError%0A%0A    def to_cpu(self):%0A        # TODO: Implement%0A        raise NotImplementedError
+self._model.cuda(device_num)%0A%0A    def to_cpu(self):%0A        self._model.cpu()
 %0A%0A  
@@ -2558,32 +2558,41 @@
 size_output(self
+, new_dim
 ):%0A        # TOD
@@ -2591,43 +2591,149 @@
    #
- TODO: Required for spaCy
+self.weight = nn.Parameter(F.pad(self.weight, ...)) # add classes%0A        #self.weight = nn.Parameter(F.pad(model.weight, ...)) #
  add 
+c
 la
-bel
+sses
 %0A   
@@ -2796,63 +2796,8 @@
 f):%0A
-        # TODO: Not required yet, but should be useful%0A
     
@@ -2923,57 +2923,195 @@
     
-# TODO: Implement%0A        raise NotImplementedError
+if self.id in params:%0A            backup = self.to_bytes()%0A            self.from_bytes(params%5Bself.id%5D)%0A        else:%0A            backup = None%0A        yield%0A        self.from_bytes(backup)
 %0A%0A

Work on PyTorch wrapper