import numpy as np
from typing import Optional, List, Tuple, Literal
from easynn.schedulers import Scheduler
import logging

logger = logging.getLogger(__name__)


# === Activation Functions ===


class ActivationFunction:
    # Values are the pre-activation values (z)
    def forward(self, values: np.ndarray) -> np.ndarray:
        raise NotImplementedError

    def backward(self, values: np.ndarray) -> np.ndarray:
        raise NotImplementedError


class ReLU(ActivationFunction):
    def forward(self, values: np.ndarray) -> np.ndarray:
        return np.maximum(values, 0)

    def backward(self, values: np.ndarray) -> np.ndarray:
        return np.where(values > 0, 1.0, 0.0)


class LeakyReLU(ActivationFunction):
    def __init__(self, leak: float = 1e-5) -> None:
        self.leak = leak
        super().__init__()

    def forward(self, values: np.ndarray) -> np.ndarray:
        return np.where(values >= 0, values, values * self.leak)

    def backward(self, values: np.ndarray) -> np.ndarray:
        return np.where(values >= 0, 1.0, self.leak)

class Linear(ActivationFunction):
    def forward(self, values: np.ndarray) -> np.ndarray:
        return values
    
    def backward(self, values: np.ndarray) -> np.ndarray:
        return np.ones_like(values)

class Softmax(ActivationFunction):
    def forward(self, values: np.ndarray) -> np.ndarray:
        exp_values = np.exp(values - np.max(values, axis=1, keepdims=True))
        return exp_values / np.sum(exp_values, axis=1, keepdims=True)

    def backward(self, values: np.ndarray) -> np.ndarray:
        s = self.forward(values)
        return s * (1 - s)
    
# === Loss Functions ===


class LossFunction:
    def forward(self, y_pred: np.ndarray, y_true: np.ndarray) -> float:
        raise NotImplementedError

    def backward(self, y_pred: np.ndarray, y_true: np.ndarray) -> np.ndarray:
        raise NotImplementedError


class MSELoss(LossFunction):
    def forward(self, y_pred: np.ndarray, y_true: np.ndarray) -> float:
        return float(np.mean((y_pred - y_true) ** 2))

    def backward(self, y_pred: np.ndarray, y_true: np.ndarray) -> np.ndarray:
        return 2 * (y_pred - y_true) / y_true.size


class CrossEntropyLoss(LossFunction):
    def __init__(self) -> None:
        super().__init__()
        self._override_activation_loss: bool = False  # Used if activation and loss are combined for efficiency; Set by FFNN if needed;
    
    def set_override_activation_loss(self, override: bool = True) -> None:
        self._override_activation_loss = override

    def forward(self, y_pred: np.ndarray, y_true: np.ndarray) -> float:
        eps = 1e-9
        y_pred = np.clip(y_pred, eps, 1 - eps)
        return -float(np.mean(np.sum(y_true * np.log(y_pred), axis=1)))
    

    def backward(self, y_pred: np.ndarray, y_true: np.ndarray) -> np.ndarray:
        if not self._override_activation_loss:
            eps = 1e-9
            y_pred = np.clip(y_pred, eps, 1 - eps)
            return - (y_true / y_pred) / y_true.shape[0]
        else:
            return (y_pred - y_true) / y_true.shape[0]

# === Regularization Functions ===
class Regularization:
    def __init__(self, reg_lambda: float, mode: Literal['l1', 'l2'] = 'l2') -> None:
        self.reg_lambda = reg_lambda
        if mode not in ['l1', 'l2']:
            raise ValueError("mode must be 'l1' or 'l2'")
        self.mode = mode

    def compute_penalty(self, weights: np.ndarray, biases: np.ndarray) -> float:
        if self.mode == 'l2':
            return self.reg_lambda * (np.sum(weights ** 2) + np.sum(biases ** 2))
        else:  # l1
            return self.reg_lambda * (np.sum(np.abs(weights)) + np.sum(np.abs(biases)))

    def compute_gradient(self, weights: np.ndarray, biases: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        if self.mode == 'l2':
            return 2 * self.reg_lambda * weights, 2 * self.reg_lambda * biases
        else:  # l1
            return self.reg_lambda * np.sign(weights), self.reg_lambda * np.sign(biases)


# === Layer Definition ===


class Layer:
    def __init__(
        self, input_dim: int, num_nodes: int, activation_function: ActivationFunction, regularization: Optional[Regularization] = None
    ) -> None:
        self.input_dim, self.num_nodes = input_dim, num_nodes

        # Initialize weights and biases identically to PyTorch's default initialization
        stdv = 1. / np.sqrt(input_dim * num_nodes)
        self.weights = np.random.rand(input_dim, num_nodes) * 2 * stdv - stdv
        self.biases = np.random.rand(1, num_nodes) * 2 * stdv - stdv
    
        self.activation_function = activation_function
        self.last_input: Optional[np.ndarray] = None
        self.last_z: Optional[np.ndarray] = None  # Pre-activation values

        self.regularization = regularization
        self._override_activation_loss: bool = False # Used if activation and loss are combined for efficiency; Set by FFNN if needed;

    def set_override_activation_loss(self, override: bool = True) -> None:
        self._override_activation_loss = override

    def inner_product(self, values: np.ndarray) -> np.ndarray:
        # (batch_size, input_dim) @ (input_dim, num_nodes) → (batch_size, num_nodes)
        return values @ self.weights + self.biases

    def forward(self, values: np.ndarray) -> np.ndarray:
        self.last_input = values
        self.last_z = self.inner_product(values)
        return self.activation_function.forward(self.last_z)

    def backward(
        self, upstream_grad: np.ndarray
    ) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        if self.last_input is None or self.last_z is None:
            raise ValueError("Must call forward before backward")

        batch_size = upstream_grad.shape[0]

        # Gradient of activation
        if self._override_activation_loss:
            activation_grad = np.ones_like(self.last_z)
        else:
            activation_grad = self.activation_function.backward(self.last_z)
        
        local_grad = upstream_grad * activation_grad

        # Gradients for weights and biases
        bias_grad = np.sum(local_grad, axis=0, keepdims=True) / batch_size
        weights_grad = (self.last_input.T @ local_grad) / batch_size
        if self.regularization is not None:
            reg_weights_grad, reg_biases_grad = self.regularization.compute_gradient(self.weights, self.biases)
            weights_grad += reg_weights_grad
            bias_grad += reg_biases_grad

        # Gradient for previous layer
        previous_grad = local_grad @ self.weights.T

        return previous_grad, (bias_grad, weights_grad)
    
    def compute_regularization_penalty(self) -> float:
        if self.regularization is not None:
            return self.regularization.compute_penalty(self.weights, self.biases)
        return 0.0


# === Feedforward Network ===


class FFNN:
    def __init__(
        self, layers: List[Layer], scheduler: Scheduler, loss_fn: LossFunction
    ) -> None:
        self.layers = layers
        self.scheduler = scheduler
        self.loss_fn = loss_fn

        if isinstance(self.loss_fn, CrossEntropyLoss):
            self.loss_fn.set_override_activation_loss(True)
            if isinstance(self.layers[-1].activation_function, Softmax):
                self.layers[-1].set_override_activation_loss(True)
                

    def predict(self, X: np.ndarray) -> np.ndarray:
        if X.ndim == 1:
            X = X.reshape(1, -1)

        values = X
        for layer in self.layers:
            values = layer.forward(values)
        return values

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        assert X.ndim == 2, f"X should be 2D, got {X.ndim}D"
        assert y.ndim == 2, f"y should be 2D, got {y.ndim}D"
        assert X.shape[0] == y.shape[0], "Batch sizes of X and y must match"

        iteration = 0
        while self.scheduler.cont:
            # Forward pass
            y_pred = self.predict(X)

            # Compute loss and gradient
            loss_value = self.loss_fn.forward(y_pred, y)
            logger.debug(f"Iteration {iteration + 1},  Pre-Regularization Loss: {loss_value}")
            # Add regularization penalties
            for layer in self.layers:
                loss_value += layer.compute_regularization_penalty()
        
            loss_grad = self.loss_fn.backward(y_pred, y)

            # Backward pass
            grads = []
            current_grad = loss_grad
            for layer in reversed(self.layers):
                current_grad, (bias_grad, weights_grad) = layer.backward(current_grad)
                grads.append((bias_grad, weights_grad))

            # Reverse grads to match layer order
            updates = self.scheduler.update(grads[::-1])

            # Apply parameter updates
            for layer, (bias_update, weights_update) in zip(self.layers, updates):
                layer.biases -= bias_update
                layer.weights -= weights_update

            iteration += 1
            logger.info(f"Iteration {iteration}, Loss: {loss_value}")
            self.scheduler.record_loss(loss_value)


# === Example usage ===

if __name__ == "__main__":
    from easynn.schedulers import GradientDescentScheduler, AdamScheduler
    from sklearn.preprocessing import StandardScaler
    x = np.linspace(0, 1, 100)
    X = np.array([x, x + np.random.randn(len(x)), x + np.random.randn(len(x))*2]).T
    X_scaler = StandardScaler().fit(X)
    X_s = X_scaler.transform(X)
    y = np.array([x + x**2 + x**3]).T
    y_scaler = StandardScaler().fit(y)
    y_s = y_scaler.transform(y)
    print(X_s)
    print(y_s)
    layers = [Layer(3, 5, ReLU(), Regularization(0.01, 'l2')), Layer(5, 25, ReLU(), Regularization(0.01, 'l2')),  Layer(25, 1, Linear())]
    network = FFNN(layers, AdamScheduler(learning_rate=0.001, epochs=1000), MSELoss())
    network.fit(X_s, y_s)
    print(np.abs(np.mean(network.predict(X_s) - y_s)))
