"""Generate paraphrased caption datasets from the existing canonical orig set.

Variants:
  - para1: Chinese paraphrase #1 (different sentence structure, same meaning)
  - para2: Chinese paraphrase #2 (different verb/wording, same meaning)
  - en: English mirror (literal translation, same meaning)

For each existing caption_*.jsonl record, the `prompt` field is replaced by
the corresponding (mode, family, variant) entry from the table below. All
other fields (raw_ts, gt_metrics, meta, ...) are preserved bit-identical.

Robustness rationale: same record, same family, same semantics — only surface
form differs. If model is robust, all 4 variants should produce similar
factual extraction scores. Score spread = prompt sensitivity.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

# (mode, family) → {variant: prompt_template_with_<ts>}
# "<ts>" will be replaced by " <ts> </ts>" at emit time
PROMPT_VARIANTS: dict[tuple[str, str], dict[str, str]] = {
    # ============== UNIVAR ==============
    ("univar", "overall"): {
        "orig":  "请对这个时间序列做深入分析：<ts>",
        "para1": "任务：分析下方时间序列，输出对其整体形态与变化脉络的综合解读。\n数据：<ts>",
        "para2": "下面这条时间序列，能不能帮我把它的整体形态和变化脉络讲清楚？<ts>",
        "en":    "Please provide an in-depth analysis of this time series: <ts>",
    },
    ("univar", "pattern"): {
        "orig":  "请分析该时间序列的主要模式及其相互关系：<ts>",
        "para1": "任务：识别下方时间序列中存在的主要模式，并刻画这些模式之间的关联结构。\n数据：<ts>",
        "para2": "看一下这条时间序列，里面主要有哪些模式？这些模式互相又是什么关系？<ts>",
        "en":    "Please identify the main patterns in this time series and discuss the interrelationships among them: <ts>",
    },
    ("univar", "stability"): {
        "orig":  "请从稳定性和可预测性的角度，对该时间序列进行分析：<ts>",
        "para1": "任务：评估下方时间序列的稳定程度，并说明它的可预测性来源。\n数据：<ts>",
        "para2": "下面这条时间序列，它稳不稳？从哪里能看出它的可预测性？<ts>",
        "en":    "Please analyze this time series from the perspectives of stability and predictability: <ts>",
    },
    ("univar", "risk"): {
        "orig":  "请分析该时间序列的主要风险和异常：<ts>",
        "para1": "任务：评估下方时间序列的结构变化风险与异常情况，结合整体趋势给出研判。\n数据：<ts>",
        "para2": "看下这条时间序列，有没有结构变化的风险或者异常？整体趋势能不能说明问题？<ts>",
        "en":    "Please identify the main risks and anomalies in this time series: <ts>",
    },
    # ============== BIVAR ==============
    ("bivar", "overall"): {
        "orig":  "请对这个双变量时间序列做深入分析：<ts>",
        "para1": "任务：分析下方双变量时间序列，输出对其整体形态与变量间关系的综合解读。\n数据：<ts>",
        "para2": "下面这对双变量时间序列，能不能帮我把它们整体的形态和相互关系讲清楚？<ts>",
        "en":    "Please provide an in-depth analysis of this bivariate time series: <ts>",
    },
    ("bivar", "pattern"): {
        "orig":  "请分析该双变量时间序列的主要关系模式及其相互作用：<ts>",
        "para1": "任务：识别下方双变量时间序列中存在的主要关系模式，并刻画这些模式之间的相互作用。\n数据：<ts>",
        "para2": "看一下这对双变量时间序列，里面主要有哪些关系模式？两条序列又是怎么互相作用的？<ts>",
        "en":    "Please analyze the main relational patterns and interactions in this bivariate time series: <ts>",
    },
    ("bivar", "stability"): {
        "orig":  "请从稳定性和共同变化的角度，对该双变量时间序列进行分析：<ts>",
        "para1": "任务：评估下方双变量时间序列的关系结构稳定程度，并说明它的可预测性来源。\n数据：<ts>",
        "para2": "下面这对双变量时间序列，它们整体的协同关系稳不稳？从哪里能看出它的可预测性？<ts>",
        "en":    "Please analyze this bivariate time series from the perspectives of stability and co-movement: <ts>",
    },
    ("bivar", "risk"): {
        "orig":  "请评估这对双变量时间序列的协同风险和异常表现：<ts>",
        "para1": "任务：评估下方双变量时间序列的协同风险与异常表现，结合整体演化轨迹给出研判。\n数据：<ts>",
        "para2": "看下这对双变量时间序列，有没有协同方面的风险或者异常表现？整体走势能不能说明问题？<ts>",
        "en":    "Please assess the joint risks and anomalous behaviors of this bivariate time series: <ts>",
    },
    # ============== MULTIVAR ==============
    ("multivar", "overall"): {
        "orig":  "请对这个多变量时间序列系统做深入分析：<ts>",
        "para1": "任务：分析下方多变量时间序列系统，输出对其整体结构与动态模式的综合解读。\n数据：<ts>",
        "para2": "下面这套多变量时间序列，能不能帮我把整体的结构和动态变化模式讲清楚？<ts>",
        "en":    "Please provide an in-depth analysis of this multivariate time series system: <ts>",
    },
    ("multivar", "pattern"): {
        "orig":  "请分析该多变量时间序列系统的主要模式特征及其相互关系：<ts>",
        "para1": "任务：识别下方多变量时间序列系统中存在的主要模式，并刻画这些模式之间的关联结构。\n数据：<ts>",
        "para2": "看一下这套多变量时间序列，里面主要有哪些模式？这些模式互相又是什么关系？<ts>",
        "en":    "Please analyze the main pattern features and their interrelationships in this multivariate time series system: <ts>",
    },
    ("multivar", "stability"): {
        "orig":  "请判断该多变量系统的协同结构是否稳定，并说明其可预测性的来源：<ts>",
        "para1": "任务：评估下方多变量时间序列系统的协同结构稳定程度，并说明它的可预测性来源。\n数据：<ts>",
        "para2": "下面这套多变量时间序列，它的整体协同结构稳不稳？从哪里能看出它的可预测性？<ts>",
        "en":    "Please determine whether the coordinated structure of this multivariate system is stable, and explain the source of its predictability: <ts>",
    },
    ("multivar", "risk"): {
        "orig":  "请评估该多变量时间序列的结构脆弱性与异常风险，并结合整体变化给出分析：<ts>",
        "para1": "任务：评估下方多变量时间序列系统的结构脆弱性与异常风险，结合整体演化轨迹给出研判。\n数据：<ts>",
        "para2": "看下这套多变量时间序列，有没有结构上的脆弱点或者异常风险？整体走势能不能说明问题？<ts>",
        "en":    "Please evaluate the structural fragility and anomaly risks of this multivariate time series, and analyze in light of overall changes: <ts>",
    },
}


def _swap_prompt(template_with_ts: str) -> str:
    """Expand <ts> → ' <ts> </ts>' to match the format used by run_ts_align."""
    return template_with_ts.replace("<ts>", " <ts> </ts>")


def build_variant(src_path: Path, dst_path: Path, variant: str) -> tuple[int, int]:
    """Read records from src, replace prompt with variant, write to dst."""
    n = 0
    n_missing = 0
    with src_path.open() as fin, dst_path.open("w") as fout:
        for line in fin:
            line = line.strip()
            if not line:
                continue
            r = json.loads(line)
            meta = r.get("meta") or {}
            mode = meta.get("prompt_mode")
            family = r.get("prompt_family")
            entry = PROMPT_VARIANTS.get((mode, family))
            if entry is None or variant not in entry:
                n_missing += 1
                continue
            r = dict(r)  # shallow copy
            r["prompt"] = _swap_prompt(entry[variant])
            r["prompt_variant"] = variant
            r["request_id"] = f"{r['id']}:caption:{variant}"
            fout.write(json.dumps(r, ensure_ascii=False) + "\n")
            n += 1
    return n, n_missing


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src-mvonly", required=False, default=None,
                    help="Original caption_500ts_mvonly.jsonl")
    ap.add_argument("--src-3split", required=True,
                    help="Original caption_501ts_3split.jsonl")
    ap.add_argument("--out-dir", required=True,
                    help="Where to write variant jsonls")
    ap.add_argument("--variants", default="para1,para2,en",
                    help="Comma-separated variant names (subset of orig/para1/para2/en)")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    variants = args.variants.split(",")
    for var in variants:
        if var == "orig":
            print(f"# variant 'orig' uses src files as-is; skip generation")
            continue
        for src_path_str, tag in [
            (args.src_3split, "3split_501ts"),
        ]:
            src = Path(src_path_str)
            dst = out_dir / f"caption_{tag}_{var}.jsonl"
            n, miss = build_variant(src, dst, var)
            print(f"# wrote {dst}: {n} records ({miss} missing variant)")


if __name__ == "__main__":
    main()