"""
Validate generated CoT data quality.

Usage:
    python validate_cot.py --cot_path /path/to/cot.jsonl --num_samples 10
"""

import json
import argparse
import numpy as np
from collections import Counter


def load_cot_data(path: str) -> list:
    data = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return data


def basic_stats(data: list):
    print(f"\n{'='*60}")
    print(f"  Basic Statistics")
    print(f"{'='*60}")
    print(f"Total records: {len(data)}")

    cot_lengths = [len(item["cot"]) for item in data]
    print(f"\nCoT Length (chars):")
    print(f"  Mean: {np.mean(cot_lengths):.1f}")
    print(f"  Median: {np.median(cot_lengths):.1f}")
    print(f"  Min: {np.min(cot_lengths)}")
    print(f"  Max: {np.max(cot_lengths)}")
    print(f"  Std: {np.std(cot_lengths):.1f}")

    short_threshold = 30
    short_count = sum(1 for l in cot_lengths if l < short_threshold)
    print(f"\n  Short CoTs (<{short_threshold} chars): {short_count} ({short_count/len(data)*100:.2f}%)")

    unique_rns = set(item["rn"] for item in data)
    print(f"\nUnique query records covered: {len(unique_rns)}")

    rn_counts = Counter(item["rn"] for item in data)
    print(f"Instructions per record: mean={np.mean(list(rn_counts.values())):.1f}")


def show_samples(data: list, num_samples: int = 5):
    print(f"\n{'='*60}")
    print(f"  Random Samples")
    print(f"{'='*60}")

    indices = np.random.choice(len(data), min(num_samples, len(data)), replace=False)
    for idx in indices:
        item = data[idx]
        print(f"\n--- Sample (rn={item['rn'][:8]}..., idx={item['prompt_index']}) ---")
        print(f"Instruction: {item['instruction']}")
        print(f"CoT: {item['cot']}")
        print()


def main(args):
    print(f"Loading CoT data from: {args.cot_path}")
    data = load_cot_data(args.cot_path)

    if not data:
        print("No data loaded!")
        return

    basic_stats(data)
    show_samples(data, args.num_samples)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Validate CoT data quality")
    parser.add_argument("--cot_path", type=str, required=True)
    parser.add_argument("--num_samples", type=int, default=5)
    args = parser.parse_args()

    main(args)