"""The model: its parameters, its prediction and one training step."""

from dataclasses import dataclass

from data import Sample


@dataclass(frozen=True)
class Model:
    """One weight that predicts `weight * x`. Replace with your own model."""

    weight: float = 0.0

    def predict(self, x: float) -> float:
        """The model's output for one input."""
        return self.weight * x

    def loss(self, rows: list[Sample]) -> float:
        """Mean squared error over the rows."""
        return sum((self.predict(row.x) - row.y) ** 2 for row in rows) / len(rows)

    def step(self, rows: list[Sample], learning_rate: float) -> "Model":
        """The model after one gradient step."""
        gradient = 2.0 * sum(row.x * (self.predict(row.x) - row.y) for row in rows) / len(rows)
        return Model(self.weight - learning_rate * gradient)
