import numpy as np
from typing import List, Tuple, Optional
import logging

logger = logging.getLogger(__name__)

gradient_type = List[Tuple[np.ndarray, np.ndarray]]


class Scheduler:
    def __init__(self) -> None:
        self._loss_history: List[float] = []

    @property
    def cont(self) -> bool:
        return False
    
    def record_loss(self, loss: float) -> None:
        self._loss_history.append(loss)
    
    @property
    def loss_history(self) -> np.ndarray:
        return np.array(self._loss_history)

    def update(self, gradients: gradient_type) -> gradient_type:
        raise NotImplementedError


class BasicScheduler(Scheduler):
    def __init__(self, learning_rate: float = 0.01, epochs: int = 100) -> None:
        super().__init__()
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.current_epoch = 0

    @property
    def cont(self) -> bool:
        return self.current_epoch < self.epochs


class GradientDescentScheduler(BasicScheduler):
    def update(self, gradients: gradient_type) -> gradient_type:
        self.current_epoch += 1
        return [
            (weight_grad * self.learning_rate, bias_grad * self.learning_rate)
            for weight_grad, bias_grad in gradients
        ]


class AdvancedScheduler(BasicScheduler):
    def _pre_update(self, gradients: gradient_type) -> None:
        raise NotImplementedError  # Defined in specific optimizers

    def _single_update(
        self, new_grad: np.ndarray, list_index: int, tuple_index: int
    ) -> np.ndarray:
        raise NotImplementedError  # Defined in specific optimizers

    def _post_update(self, gradients: gradient_type, updates: gradient_type) -> None:
        raise NotImplementedError  # Defined in specific optimizers

    def update(self, gradients: gradient_type) -> gradient_type:
        logger.debug(f"Epoch {self.current_epoch + 1}: Updating gradients.")
        self.current_epoch += 1
        self._pre_update(gradients)

        updates = [
            (
                self._single_update(new_w_grad, i, 0),
                self._single_update(new_b_grad, i, 1),
            )
            for i, (new_w_grad, new_b_grad) in enumerate(gradients)
        ]
        logger.debug(f"Updates computed. Total update sum: {sum(np.sum(u) for pair in updates for u in pair)}")

        self._post_update(gradients, updates)
        return updates


class MomentumScheduler(AdvancedScheduler):
    def __init__(self, *args, delta: float = 1.0, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.delta = delta
        self.prev_updates: Optional[gradient_type] = None

    def _pre_update(self, gradients: gradient_type) -> None:
        if self.prev_updates is None:
            self.prev_updates = [
                (np.zeros_like(w), np.zeros_like(b)) for w, b in gradients
            ]

    def _single_update(
        self, new_grad: np.ndarray, list_index: int, tuple_index: int
    ) -> np.ndarray:
        assert self.prev_updates is not None
        old_grad = self.prev_updates[list_index][tuple_index]
        return self.delta * old_grad + (1 - self.delta) * new_grad * self.learning_rate

    def _post_update(self, gradients: gradient_type, updates: gradient_type) -> None:
        self.prev_updates = updates


class AdaGradScheduler(AdvancedScheduler):
    def __init__(self, *args, epsilon: float = 1e-8, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.epsilon = epsilon
        self.quadratic_sum: Optional[gradient_type] = None

    def _pre_update(self, gradients: gradient_type) -> None:
        if self.quadratic_sum is None:
            self.quadratic_sum = [
                (np.zeros_like(w), np.zeros_like(b)) for w, b in gradients
            ]
        else:
            self.quadratic_sum = [
                (
                    old_w_sum + new_w_grad**2,
                    old_b_sum + new_b_grad**2,
                )
                for (old_w_sum, old_b_sum), (new_w_grad, new_b_grad) in zip(
                    self.quadratic_sum, gradients
                )
            ]

    def _single_update(
        self, new_grad: np.ndarray, list_index: int, tuple_index: int
    ) -> np.ndarray:
        assert self.quadratic_sum is not None
        quad_sum = self.quadratic_sum[list_index][tuple_index]
        return new_grad / (np.sqrt(quad_sum) + self.epsilon) * self.learning_rate

    def _post_update(self, gradients: gradient_type, updates: gradient_type) -> None:
        pass


class RMSPropScheduler(AdvancedScheduler):
    def __init__(
        self, *args, gamma: float = 0.9, epsilon: float = 1e-8, **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.gamma = gamma
        self.epsilon = epsilon
        self.running_avg: Optional[gradient_type] = None

    def _pre_update(self, gradients: gradient_type) -> None:
        if self.running_avg is None:
            self.running_avg = [
                (np.zeros_like(w), np.zeros_like(b)) for w, b in gradients
            ]
        else:
            self.running_avg = [
                (
                    self.gamma * old_w_avg + (1 - self.gamma) * (new_w_grad**2),
                    self.gamma * old_b_avg + (1 - self.gamma) * (new_b_grad**2),
                )
                for (old_w_avg, old_b_avg), (new_w_grad, new_b_grad) in zip(
                    self.running_avg, gradients
                )
            ]

    def _single_update(
        self, new_grad: np.ndarray, list_index: int, tuple_index: int
    ) -> np.ndarray:
        assert self.running_avg is not None
        run_avg = self.running_avg[list_index][tuple_index]
        return new_grad / (np.sqrt(run_avg) + self.epsilon) * self.learning_rate

    def _post_update(self, gradients: gradient_type, updates: gradient_type) -> None:
        pass


class AdamScheduler(AdvancedScheduler):
    def __init__(
        self,
        *args,
        beta1: float = 0.9,
        beta2: float = 0.999,
        epsilon: float = 1e-8,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.m: Optional[gradient_type] = None
        self.v: Optional[gradient_type] = None

    def _pre_update(self, gradients: gradient_type) -> None:
        if self.m is None or self.v is None:
            self.m = [(np.zeros_like(w), np.zeros_like(b)) for w, b in gradients]
            self.v = [(np.zeros_like(w), np.zeros_like(b)) for w, b in gradients]
        else:
            self.m = [
                (
                    self.beta1 * old_w_m + (1 - self.beta1) * new_w_grad,
                    self.beta1 * old_b_m + (1 - self.beta1) * new_b_grad,
                )
                for (old_w_m, old_b_m), (new_w_grad, new_b_grad) in zip(
                    self.m, gradients
                )
            ]
            self.v = [
                (
                    self.beta2 * old_w_v + (1 - self.beta2) * (new_w_grad**2),
                    self.beta2 * old_b_v + (1 - self.beta2) * (new_b_grad**2),
                )
                for (old_w_v, old_b_v), (new_w_grad, new_b_grad) in zip(
                    self.v, gradients
                )
            ]

    def _single_update(
        self, new_grad: np.ndarray, list_index: int, tuple_index: int
    ) -> np.ndarray:
        assert self.m is not None and self.v is not None
        m = self.m[list_index][tuple_index]
        v = self.v[list_index][tuple_index]
        m_hat = m / (1 - self.beta1**self.current_epoch)
        v_hat = v / (1 - self.beta2**self.current_epoch)
        return m_hat / (np.sqrt(v_hat) + self.epsilon) * self.learning_rate

    def _post_update(self, gradients: gradient_type, updates: gradient_type) -> None:
        pass
