import numpy as np
import pandas as pd

from timesfm_app.core import ForecastConfig, ForecastPipeline, TimesFMService


class FakeTimesFMService(TimesFMService):
    """
    Replace the model call with deterministic values for tests.
    Returns predictions: [1, 2, ..., max_h] for each series in the batch.
    """

    def __init__(self):
        pass

    def forecast_batch(self, past_values, freq_code: int, max_h: int) -> np.ndarray:
        base = np.arange(1, max_h + 1, dtype=float)
        return np.tile(base, (len(past_values), 1))


def test_pipeline_long_output_shape():
    df = pd.DataFrame(
        {
            "date": pd.date_range("2023-01-01", periods=5, freq="ME").strftime(
                "%Y-%m-%d"
            ),
            "value": [1, 2, 3, 4, 5],
            "key": ["A"] * 5,
        }
    )

    cfg = ForecastConfig(
        model_name="dummy",
        device="cpu",
        batch_size=2,
        timesfm_freq_code=2,
        pandas_freq_fallback="M",
        horizon_min=1,
        horizon_max=3,
    )

    svc = FakeTimesFMService()
    pipe = ForecastPipeline(svc=svc, config=cfg)
    out = pipe.run(df)

    assert len(out) == 3
    assert out["key"].unique().tolist() == ["A"]
    assert out["step"].tolist() == [1, 2, 3]
    assert np.allclose(out["forecast"].values, [1.0, 2.0, 3.0])