"""A Model class holds the top node of the whole structure.
"""
import numpy as np


class SequentialModel:
    def __init__(self, cost_func: callable):
        """A sequential model is a type of model which the data
        flows layer by layer, without skipping or having multiple 
        source of inputs or slots of outputs.

        Args:
            cost_func (callable): Cost function
        """        
        self.layers = []
        self.cost_func = cost_func

    def train(self, x: np.ndarray, y: np.ndarray, n_iters: int, lr: float):
        """Train the model

        Args:
            x (np.ndarray): Input data
            y (np.ndarray): Ground truth
            n_iters (int): Number of iterations
            lr (float): Learning rate
        """             
        self.predict(x)
        weights, dl_dz, dl_dw, dl_db = None, None, None, None
        m = x.shape[0]
        for it in range(n_iters):
            loss = self.cost_func(self.predict(x), y)
            print(f"Iteration {it}# train loss: {loss}")

            yhat = self.predict(x)
            for i in reversed(range(len(self.layers))):
                if i == len(self.layers) - 1:
                    # For the last layer, it has to get the gradients from
                    # the loss, while other layers get the gradients from the 
                    # layer after.
                    dl_dz = self.cost_func(y, yhat, prime=True) * self.layers[i].activation(self.layers[i].latest_output, prime=True) / m
                    dl_dw = dl_dz.T @ self.layers[i - 1].latest_output
                    dl_db = np.sum(dl_dz, axis=0, keepdims=True).T
                else:
                    weights = self.layers[i + 1].weights
                    dl_dz = dl_dz @ weights * self.layers[i].activation(self.layers[i].latest_output, prime=True) / m
                    # If layer[i - 1] is the input layer, then dl_dw should be the input
                    dl_dw = dl_dz.T @ self.layers[i - 1].latest_output if i > 0 else dl_dz.T @ x
                    dl_db = np.sum(dl_dz, axis=0, keepdims=True).T
                self.layers[i].dl_dw = dl_dw
                self.layers[i].dl_db = dl_db

            for layer in self.layers:
                layer.update(lr)

    def append(self, *args):
        """Append one or more layers to the sequential model
        """        
        for arg in args:
            self.layers.append(arg)

    def predict(self, data: np.ndarray) -> np.ndarray:
        """Forward pass through all the layers and get the
        predicted output

        Args:
            data (np.ndarray): Input data

        Returns:
            np.ndarray: Predicted output.
        """        
        for layer in self.layers:
            data = layer(data)
        return data
