lsm_regressor_clean

from typing import Sequence, Union
from numbers import Real

from regressor_abc import RegressorABC


class RegressorLSM(RegressorABC):
    _a: int
    _b: int

    def fit(self, abscissa: Sequence[Real], ordinates: Sequence[Real]) -> None:
        x_avg = 0
        y_avg = 0
        xy_avg = 0
        x2_avg = 0
        n = len(abscissa)
        for i in range(n):
            x_avg += abscissa[i]
            y_avg += ordinates[i]
            xy_avg += abscissa[i] * ordinates[i]
            x2_avg += abscissa[i] ** 2
        x_avg /= n
        y_avg /= n
        xy_avg /= n
        x2_avg /= n
        self._a = (xy_avg - x_avg * y_avg) / (x2_avg - x_avg ** 2)
        self._b = y_avg - self._a * x_avg

    def predict(self, abscissa: Union[Real, Sequence[Real]]) -> list:
        if isinstance(abscissa, Real):
            return [self._a * abscissa + self._b]
        return [(self._a * x + self._b) for x in abscissa]