#!/usr/bin/env python3
"""Export task-level model input CSV files from normalized TS-Bench tables.

This script does not fetch data. It reads:
  normalized/observations.csv
  normalized/forecast_tasks.csv

For each task it writes:
  model_inputs/<task_id>/context.csv
  model_inputs/<task_id>/future_target.csv

The context file is what a model can consume. The future_target file is only
for evaluator/scoring code.
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple


UTC = timezone.utc


def parse_dt(value: str) -> datetime:
    text = value.strip()
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    dt = datetime.fromisoformat(text)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=UTC)
    return dt.astimezone(UTC)


def read_csv(path: Path) -> List[Dict[str, str]]:
    with path.open("r", newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def write_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: List[str]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow({name: row.get(name, "") for name in fieldnames})
            count += 1
    return count


def pivot_rows(rows: List[Dict[str, str]]) -> Tuple[List[Dict[str, Any]], List[str]]:
    grouped: Dict[str, Dict[str, Any]] = {}
    variables = sorted({row["variable_id"] for row in rows})
    for row in rows:
        timestamp = row["timestamp"]
        if timestamp not in grouped:
            grouped[timestamp] = {"timestamp": timestamp}
        grouped[timestamp][row["variable_id"]] = row["value"]
    return [grouped[key] for key in sorted(grouped)], ["timestamp"] + variables


def export_task(
    task: Dict[str, str],
    observations: List[Dict[str, str]],
    output_root: Path,
) -> Dict[str, Any]:
    dataset_id = task["dataset_id"]
    entity_id = task["target_entity_id"]
    target_variable_id = task["target_variable_id"]
    forecast_issue_time = parse_dt(task["forecast_issue_time"])
    context_start = parse_dt(task["context_start"])
    context_end = parse_dt(task["context_end"])
    horizon_start = parse_dt(task["horizon_start"])
    horizon_end = parse_dt(task["horizon_end"])

    context_rows = []
    future_rows = []
    for row in observations:
        if row["dataset_id"] != dataset_id or row["entity_id"] != entity_id:
            continue
        timestamp = parse_dt(row["timestamp"])
        available_time = parse_dt(row["available_time"])
        if context_start <= timestamp <= context_end and available_time <= forecast_issue_time:
            context_rows.append(row)
        if horizon_start <= timestamp <= horizon_end and row["variable_id"] == target_variable_id:
            future_rows.append(row)

    context_pivot, context_fields = pivot_rows(context_rows)
    future_pivot, future_fields = pivot_rows(future_rows)

    task_dir = output_root / task["task_id"]
    context_count = write_csv(task_dir / "context.csv", context_pivot, context_fields)
    future_count = write_csv(task_dir / "future_target.csv", future_pivot, future_fields)
    (task_dir / "task.json").write_text(json.dumps(task, ensure_ascii=False, indent=2), encoding="utf-8")

    return {
        "task_id": task["task_id"],
        "dataset_id": dataset_id,
        "entity_id": entity_id,
        "target_variable_id": target_variable_id,
        "context_rows": context_count,
        "future_target_rows": future_count,
        "context_csv": str(task_dir / "context.csv"),
        "future_target_csv": str(task_dir / "future_target.csv"),
    }


def parse_args(argv: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export model input CSVs from normalized TS-Bench tables.")
    parser.add_argument("--input", default="data/sample_week/normalized", help="Directory containing normalized CSV tables.")
    parser.add_argument("--output", default="data/sample_week/model_inputs", help="Output directory for task-level model inputs.")
    parser.add_argument("--limit", type=int, default=20, help="Maximum number of tasks to export.")
    return parser.parse_args(argv)


def main(argv: List[str]) -> int:
    args = parse_args(argv)
    input_dir = Path(args.input)
    output_dir = Path(args.output)
    observations_path = input_dir / "observations.csv"
    tasks_path = input_dir / "forecast_tasks.csv"
    if not observations_path.exists() or not tasks_path.exists():
        raise FileNotFoundError(f"Expected {observations_path} and {tasks_path}")

    observations = read_csv(observations_path)
    tasks = read_csv(tasks_path)[: args.limit]
    manifest = [export_task(task, observations, output_dir) for task in tasks]
    write_csv(
        output_dir / "manifest.csv",
        manifest,
        [
            "task_id",
            "dataset_id",
            "entity_id",
            "target_variable_id",
            "context_rows",
            "future_target_rows",
            "context_csv",
            "future_target_csv",
        ],
    )
    (output_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Exported {len(manifest)} tasks to {output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))