import torch

import torch.nn as nn

from torch.nn.parameter import Parameter

class LinearRegressionModel(nn.Module):

    def __init__(self, p, weights = None, bias = None):

        super(LinearRegressionModel, self).__init__()

        self.linear = nn.Linear(p, 1)

        if weights is not None:

            self.linear.weight = Parameter(torch.Tensor([weights]))

        if bias is not None:

            self.linear.bias = Parameter(torch.Tensor([bias]))

    def forward(self, x):

        return self.linear(x)

class LogisticRegressionModel(nn.Module):

    def __init__(self, p, weights = None, bias = None):

        super(LogisticRegressionModel, self).__init__()

        self.linear = nn.Linear(p, 1)

        if weights is not None:

            self.linear.weight = Parameter(torch.Tensor([weights]))

        if bias is not None:

            self.linear.bias = Parameter(torch.Tensor([bias]))

    def forward(self, x):

        return torch.sigmoid(self.linear(x))

# model_modules["Logistic"](3, (1,1,1), 0).forward(torch.zeros([1,3]))