#!/usr/bin/env python3
"""Train a tiny demo weather model and export it to ONNX.

This is intentionally simple: it learns from daily max temperature only and
predicts the next five max-temperature values from the previous seven days.
Use it as a scaffold for replacing the mock model with a stronger architecture.
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim


LOOKBACK_DAYS = 7
FORECAST_DAYS = 5
HISTORY_CSV_PATH = Path("data/tokyo_weather_history_30y.csv")
ONNX_PATH = Path("models/mock_model/model.onnx")
PT_PATH = Path("models/mock_model/model.pt")


class MovingAverageForecastModel(nn.Module):
    def __init__(self, lookback: int, forecast: int) -> None:
        super().__init__()
        self.fc = nn.Linear(lookback, forecast)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc(x)


def create_sequences(values: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
    x_values = []
    y_values = []
    for index in range(len(values) - LOOKBACK_DAYS - FORECAST_DAYS + 1):
        x_values.append(values[index : index + LOOKBACK_DAYS])
        y_values.append(values[index + LOOKBACK_DAYS : index + LOOKBACK_DAYS + FORECAST_DAYS])
    return (
        torch.tensor(np.array(x_values), dtype=torch.float32),
        torch.tensor(np.array(y_values), dtype=torch.float32),
    )


def load_training_values() -> np.ndarray:
    if not HISTORY_CSV_PATH.exists():
        raise FileNotFoundError(
            f"{HISTORY_CSV_PATH} does not exist. Run scripts/fetch_historical_weather.py first."
        )
    history = pd.read_csv(HISTORY_CSV_PATH)
    return history["high_c"].dropna().to_numpy(dtype=np.float32)


def main() -> None:
    values = load_training_values()
    x_train, y_train = create_sequences(values)

    model = MovingAverageForecastModel(LOOKBACK_DAYS, FORECAST_DAYS)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    with torch.no_grad():
        model.fc.weight.fill_(1.0 / LOOKBACK_DAYS)
        model.fc.bias.fill_(0.0)

    model.train()
    for epoch in range(120):
        optimizer.zero_grad()
        predictions = model(x_train)
        loss = criterion(predictions, y_train)
        loss.backward()
        optimizer.step()

    ONNX_PATH.parent.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), PT_PATH)
    dummy_input = torch.zeros(1, LOOKBACK_DAYS, dtype=torch.float32)
    torch.onnx.export(
        model,
        dummy_input,
        ONNX_PATH,
        input_names=["input"],
        output_names=["forecast"],
        dynamic_axes={"input": {0: "batch"}, "forecast": {0: "batch"}},
        opset_version=17,
    )
    print(f"Saved PyTorch weights to {PT_PATH}")
    print(f"Exported ONNX model to {ONNX_PATH}")


if __name__ == "__main__":
    main()