"""OpenDots example. Default mode checks baselines without model dependencies.

The optional model call uses Mistral's documented Timestral API.
Upstream code: https://github.com/mistralai/timestral (Apache-2.0).
Model weights have a separate Modified MIT license.
"""
import argparse
import json
import math
from pathlib import Path

MODEL_ID = "mistralai/Timestral-5M"
MODEL_REVISION = "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f"


def mae(actual, predicted):
    if not actual or len(actual) != len(predicted):
        raise ValueError("Use nonempty, equal-length actual and predicted values.")
    if not all(math.isfinite(float(x)) for x in [*actual, *predicted]):
        raise ValueError("This example requires finite numbers.")
    return sum(abs(a - b) for a, b in zip(actual, predicted)) / len(actual)


def baselines(context, horizon):
    if len(context) < 7 or horizon < 1:
        raise ValueError("Provide at least seven context values and a positive horizon.")
    return {
        "repeat_last": [context[-1]] * horizon,
        "repeat_week": [context[-7 + (step % 7)] for step in range(horizon)],
    }


def compare(run_model=False):
    # Invented daily values with an exactly repeating seven-day pattern.
    series = [100.0 + 2.0 * (day % 7) for day in range(63)]
    context, holdout = series[:-7], series[-7:]
    predictions = baselines(context, len(holdout))
    result = {
        "data": "synthetic; exactly repeating weekly pattern",
        "context_steps": len(context),
        "holdout_steps": len(holdout),
        "model_executed": False,
        "scores": {name: {"mae": mae(holdout, values)} for name, values in predictions.items()},
    }
    if run_model:
        import numpy as np
        import torch
        from timestral import TimestralForecaster

        torch.set_num_threads(1)
        model = TimestralForecaster.from_pretrained(
            MODEL_ID, revision=MODEL_REVISION, device="cpu"
        )
        forecast = model.forecast(np.asarray(context, dtype=np.float32), prediction_length=len(holdout))
        predicted = np.asarray(forecast.median).reshape(-1).tolist()
        result["scores"][MODEL_ID] = {"mae": mae(holdout, predicted)}
        result.update(model_executed=True, model_id=MODEL_ID, model_revision=MODEL_REVISION)
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-model", action="store_true", help="Download public model files and run CPU inference.")
    parser.add_argument("--output", type=Path, default=Path("timestral-result.json"))
    args = parser.parse_args()
    if args.output.exists():
        parser.error("Output exists. Read it or choose a new --output filename.")
    result = compare(run_model=args.run_model)
    with args.output.open("x") as handle:
        json.dump(result, handle, indent=2, allow_nan=False)
        handle.write("\n")
    print(json.dumps(result, indent=2, allow_nan=False))


if __name__ == "__main__":
    main()