import numpy as np

class SingleLayerLinearRegression:
    def __init__(self):
        base_array = np.arange(10)
        self.W = np.tile(base_array, (2, 1)).reshape(-1, 1)
        self.b = np.zeros(1)

    def forward(self, x):
        return np.dot(x, self.W) + self.b

    def backward(self, x, y, output, learning_rate):
        output_error = output - y
        output_delta = output_error

        self.W -= learning_rate * np.dot(x.T, output_delta)
        self.b -= learning_rate * np.sum(output_delta, axis=0)

    def train(self, X, y, epochs, learning_rate):
        for epoch in range(epochs):
            output = self.forward(X)
            self.backward(X, y, output, learning_rate)
            if epoch % 100 == 0:
                loss = np.mean((output - y) ** 2)
                print(f"Epoch {epoch}, Loss: {loss}")

    def predict(self, x):
        return self.forward(x)

# One-hot encode inputs
def one_hot_encode(a, b, op):
    x = np.zeros(20)
    x[a] = 1
    x[b+10] = 1
    '''
    if op == '+':
        x[10] = 1
    elif op == '-':
        x[11] = 1
    elif op == '*':
        x[12] = 1
    elif op == '/':
        x[13] = 1
    '''
    return x

# Generate training data
def generate_data(num_samples=100, max_value=10):
    data = []
    for _ in range(num_samples):
        a = np.random.randint(0, max_value)
        b = np.random.randint(0, max_value)
        op = np.random.choice(['+', '-', '*', '/'])
        if op == '+':
            result = a + b
        elif op == '-':
            result = a - b
        elif op == '*':
            result = a * b
        elif op == '/':
            if b == 0:
                continue  # Avoid division by zero
            result = a / b
        data.append((a, b, op, result))
    return data

# Prepare training data
data = generate_data()
X = np.array([one_hot_encode(a, b, op) for a, b, op, _ in data])
y = np.array([result for _, _, _, result in data]).reshape(-1, 1)

# Define and train the model
lr_model = SingleLayerLinearRegression()
#lr_model.train(X, y, epochs=1000, learning_rate=0.01)

# Example input
a = 1
b = 2
op = '+'

# One-hot encode the input
encoded_input = one_hot_encode(a, b, op).reshape(1, -1)
print("Encoded Input:")
print(encoded_input)

# Feed the input to the model
predicted_output = lr_model.predict(encoded_input)
print(f"Input: {a} {op} {b} = \nPredicted Output: {predicted_output[0][0]}")

Encoded Input:
[[0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]]
Input: 1 + 2 = 
Predicted Output: 3.0