import numpy as np
import json

class LayerConfig:
    def __init__(self, name, size, activation):
        self.name = name
        self.size = size
        self.activation = activation

class SimpleMLModel:
    def __init__(self, layer_configs, learning_rate=0.01, loss='mse'):
        self.learning_rate = learning_rate
        self.loss = loss
        self.layer_configs = layer_configs
        self.model = self._init_model()

    def _init_model(self):
        model = {}
        sizes = [self.layer_configs[0].size]  # Input layer size
        
        for config in self.layer_configs[1:]:  # Exclude input layer
            sizes.append(config.size)

        for i in range(len(sizes) - 1):
            model[f'W{i}'] = np.random.randn(sizes[i], sizes[i+1]) * 0.01
            model[f'b{i}'] = np.zeros((1, sizes[i+1]))

        return model

    def forward(self, X):
        activations = [X]
        for i, config in enumerate(self.layer_configs[1:]): # Exclude input layer
            W = self.model[f'W{i}']
            b = self.model[f'b{i}']
            X = np.dot(X, W) + b

            if config.activation == 'relu':
                X = np.maximum(0, X)
            elif config.activation == 'sigmoid':
                X = 1 / (1 + np.exp(-X))
            elif config.activation == 'tanh':
                X = np.tanh(X)

            activations.append(X)
        return activations
    
    def backward(self, activations, y_true):
        grads = {}
        dA = activations[-1] - y_true
        
        for i in reversed(range(len(self.model) // 2)):
            dZ = dA * (activations[i+1] > 0)  # ReLU backward
            grads[f'dW{i}'] = np.dot(activations[i].T, dZ) / y_true.shape[0]
            grads[f'db{i}'] = np.sum(dZ, axis=0, keepdims=True) / y_true.shape[0]
            if i > 0:
                dA = np.dot(dZ, self.model[f'W{i}'].T)
        
        return grads
    
    def update_params(self, grads):
        for i in range(len(self.model) // 2):
            self.model[f'W{i}'] -= self.learning_rate * grads[f'dW{i}']
            self.model[f'b{i}'] -= self.learning_rate * grads[f'db{i}']
    
    def train(self, X, y, epochs=100):
        for epoch in range(epochs):
            activations = self.forward(X)
            grads = self.backward(activations, y)
            self.update_params(grads)
    
    def predict(self, X):
        activations = self.forward(X)
        return activations[-1]
    
    def save_model(self, filepath):
        np.savez(filepath, **self.model)
    
    def load_model(self, filepath):
        data = np.load(filepath)
        self.model = {k: data[k] for k in data}
        
    def save_config(self, filepath):
        config_list = []
        for config in self.layer_configs:
            config_list.append({
                "name": config.name,
                "size": config.size,
                "activation": config.activation
            })
        
        with open(filepath, 'w') as f:
            json.dump(config_list, f, indent=4)

    def load_config(self, filepath):
        with open(filepath, 'r') as f:
            config_list = json.load(f)

        self.layer_configs = []
        for config_data in config_list:
            self.layer_configs.append(LayerConfig(**config_data))

        self.model = self._init_model()  # Re-initialize model based on loaded config

input_size = 2 
output_size = 1

# Example: Specific floating-point values for X and y
X = np.array([
    [10, 10],
    [5, 5],
    [15, 15],
], dtype=np.float32)  # Specify dtype if needed

y = np.array([
    [20],
    [10],
    [30],
], dtype=np.float32)

# Define your model architecture using LayerConfig
layer_configs = [
    LayerConfig("input", input_size, None), 
    LayerConfig("hidden1", 16, "sigmoid"),
    #LayerConfig("hidden2", 32, "relu"),
    LayerConfig("output", output_size, None)
]

# Create and train the model
model = SimpleMLModel(layer_configs, learning_rate=0.01, loss='mse') 
model.train(X, y, epochs=1000)

# Save the trained model (optional)
model.save_model("trained_model.npz")

# Make predictions
predictions = model.predict(X)

print(predictions)