from ..operation import PyTorchOperation


class Flatten(PyTorchOperation):
    def to_init_code(self, field):
        return ''

    def to_forward_code(self, field, output, *inputs) -> str:
        assert len(inputs) == 1
        return f'{output} = {inputs[0]}.view({inputs[0]}.size(0), -1)'


class Dense(PyTorchOperation):
    def to_init_code(self, field):
        return f"self.{field} = nn.Linear({self.parameters['in_features']}, {self.parameters['out_features']})"

    def to_forward_code(self, field, output, *inputs) -> str:
        assert len(inputs) == 1
        return f'{output} = self.{field}({inputs[0]})'


class Softmax(PyTorchOperation):
    def to_init_code(self, field):
        return ''

    def to_forward_code(self, field, output, *inputs) -> str:
        assert len(inputs) == 1
        return f'{output} = F.softmax({inputs[0]}, -1)'
