from pathlib import Path
import json
import pandas as pd
from tsfm_bench.data.registry import load_data_source

source_config = Path("configs/datasets/ts_bench.yaml")
source = load_data_source(source_config)

# Map leaderboard name to task name
lb_to_task = {}
for task_name in source.list_datasets():
    lb_name = source.leaderboard_dataset_name(task_name)
    lb_to_task[lb_name] = task_name

# Now walk space/results/ and find all forecast JSONs
results_root = Path("space/results")
patched_count = 0
for p in results_root.glob("**/forecasts/*.json"):
    try:
        data = json.loads(p.read_text())
        lb_dataset = data.get("dataset")
        if lb_dataset in lb_to_task:
            task_name = lb_to_task[lb_dataset]
            record = next(iter(source.stream(task_name)))
            n_context = data.get("context_length")
            n_future = data.get("prediction_length")

            start_period = pd.Period(record.start, freq=record.freq)
            dr = pd.period_range(start=start_period, periods=n_context + n_future, freq=record.freq)
            timestamps = [str(p.to_timestamp()) for p in dr]

            data["timestamps"] = timestamps
            p.write_text(json.dumps(data, indent=2) + "\n")
            patched_count += 1
        else:
            # Fallback mapping if task names changed slightly
            # We can find task that has the dataset name as substring
            matched_task = None
            dataset_name = lb_dataset.split("/")[0]
            for t_name in source.list_datasets():
                if dataset_name in t_name:
                    matched_task = t_name
                    break
            if matched_task:
                record = next(iter(source.stream(matched_task)))
                n_context = data.get("context_length")
                n_future = data.get("prediction_length")
                start_period = pd.Period(record.start, freq=record.freq)
                dr = pd.period_range(start=start_period, periods=n_context + n_future, freq=record.freq)
                timestamps = [str(p.to_timestamp()) for p in dr]
                data["timestamps"] = timestamps
                p.write_text(json.dumps(data, indent=2) + "\n")
                patched_count += 1
            else:
                # Fallback to evaluated_at date_range
                eval_at = data.get("evaluated_at")
                freq = data.get("freq", "1h")
                n_context = data.get("context_length")
                n_future = data.get("prediction_length")
                try:
                    # Map '1h' to 'H' or 'h' if pandas complains, or handle basic units
                    p_freq = freq
                    if p_freq == "1h":
                        p_freq = "h"
                    end_dt = pd.to_datetime(eval_at).tz_convert(None)
                    dr = pd.date_range(end=end_dt, periods=n_context + n_future, freq=p_freq)
                    timestamps = [str(p) for p in dr]
                    data["timestamps"] = timestamps
                    p.write_text(json.dumps(data, indent=2) + "\n")
                    patched_count += 1
                except Exception as e:
                    print(f"Dataset {lb_dataset} not found and fallback failed: {e}")

    except Exception as e:
        print(f"Error patching {p.name}: {e}")

print(f"Successfully patched {patched_count} forecast snapshots.")