from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class BootstrapVARResult:
    intercept: np.ndarray
    coefficients: np.ndarray
    residuals: np.ndarray

    def simulate(
        self,
        initial_state: np.ndarray,
        horizon: int,
        paths: int,
        seed: int,
    ) -> np.ndarray:
        rng = np.random.default_rng(seed)
        state_dim = initial_state.shape[0]
        simulations = np.zeros((paths, horizon, state_dim), dtype=float)
        residual_idx = np.arange(self.residuals.shape[0])
        for path in range(paths):
            current = initial_state.astype(float).copy()
            for step in range(horizon):
                shock = self.residuals[rng.choice(residual_idx)]
                current = self.intercept + self.coefficients @ current + shock
                simulations[path, step, :] = current
        return simulations


class BootstrapVARModel:
    def fit(self, data: np.ndarray) -> BootstrapVARResult:
        if data.ndim != 2:
            raise ValueError("data must be 2-dimensional")
        if data.shape[0] < 3:
            raise ValueError("data must have at least 3 rows")
        x = data[:-1]
        y = data[1:]
        design = np.column_stack([np.ones(len(x)), x])
        beta, *_ = np.linalg.lstsq(design, y, rcond=None)
        fitted = design @ beta
        residuals = y - fitted
        intercept = beta[0]
        coefficients = beta[1:].T
        return BootstrapVARResult(
            intercept=intercept, coefficients=coefficients, residuals=residuals
        )
