#!/usr/bin/env python3
"""
LUNA SFT Dataset Augmentation & Purification Pipeline
=====================================================
1. Load existing cleaned SFT data
2. Deduplicate repetitive instruction prefixes (cap at 500 per prefix)
3. Download diverse, high-quality open datasets (Dolly, Alpaca, SciQ, GSM8K)
4. Clean, format, filter, and rebrand all new data
5. Add input+output only entries (no instruction) for format diversity
6. Balance underrepresented categories
7. Validate token limits (<=1024), dedup, final merge
8. Report
"""

import json, re, os, random, hashlib, collections, sys
from pathlib import Path
from transformers import AutoTokenizer

# ── Config ──
EXISTING_PATH = Path("Base/Datasets/sft_clean/all_sft_clean.json")
OUT_DIR = Path("Base/Datasets/sft_clean")
MAX_TOKENS = 1000  # leave headroom below 1024 block size
PREFIX_CAP = 500   # max entries per repeated instruction prefix
SEED = 42

random.seed(SEED)

# ── Tokenizer ──
print("Loading tokenizer...")
tok = AutoTokenizer.from_pretrained("Base/checkpoints/EleutherAI/pythia-160m")

def count_tokens(entry):
    text = entry.get("instruction", "")
    if entry.get("input", ""):
        text += "\n" + entry["input"]
    text += "\n" + entry["output"]
    return len(tok.encode(text))

def entry_hash(entry):
    key = (entry.get("instruction","").strip().lower() + "||" +
           entry.get("input","").strip().lower())
    return hashlib.md5(key.encode()).hexdigest()

# ═══════════════════════════════════════════
# STEP 1: Load existing data
# ═══════════════════════════════════════════
print("\n[1/8] Loading existing dataset...")
with open(EXISTING_PATH, "r", encoding="utf-8") as f:
    existing = json.load(f)
print(f"  Existing: {len(existing):,}")

# ═══════════════════════════════════════════
# STEP 2: Cap repetitive instruction prefixes
# ═══════════════════════════════════════════
print("\n[2/8] Capping repetitive instruction prefixes...")

prefix_groups = collections.defaultdict(list)
for i, entry in enumerate(existing):
    prefix = entry["instruction"].strip()[:50].lower()
    prefix_groups[prefix].append(i)

capped = set()
removed_for_repetition = 0
for prefix, indices in prefix_groups.items():
    if len(indices) > PREFIX_CAP:
        keep = set(random.sample(indices, PREFIX_CAP))
        for idx in indices:
            if idx not in keep:
                capped.add(idx)
                removed_for_repetition += 1

existing = [e for i, e in enumerate(existing) if i not in capped]
print(f"  Removed {removed_for_repetition:,} over-represented entries")
print(f"  Remaining: {len(existing):,}")

# Build hash set of existing entries for dedup
existing_hashes = {entry_hash(e) for e in existing}

# ═══════════════════════════════════════════
# STEP 3: Download diverse open datasets
# ═══════════════════════════════════════════
print("\n[3/8] Downloading diverse open datasets...")

from datasets import load_dataset

new_entries = []

# ── 3a: Databricks Dolly 15K (human-written, diverse, CC-licensed) ──
print("  Downloading databricks-dolly-15k...")
try:
    dolly = load_dataset("databricks/databricks-dolly-15k", split="train")
    for row in dolly:
        inst = (row.get("instruction") or row.get("context") or "").strip()
        inp = (row.get("context") or "").strip()
        out = (row.get("response") or "").strip()
        # Dolly uses instruction and context
        if row.get("instruction","").strip() and out:
            entry = {"instruction": row["instruction"].strip(), "input": inp, "output": out}
            new_entries.append(("dolly", entry))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='dolly'):,} from Dolly")
except Exception as e:
    print(f"    Dolly failed: {e}")

# ── 3b: Alpaca Cleaned (diverse instructions) ──
print("  Downloading alpaca-cleaned...")
try:
    alpaca = load_dataset("yahma/alpaca-cleaned", split="train")
    for row in alpaca:
        inst = (row.get("instruction") or "").strip()
        inp = (row.get("input") or "").strip()
        out = (row.get("output") or "").strip()
        if inst and out:
            new_entries.append(("alpaca", {"instruction": inst, "input": inp, "output": out}))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='alpaca'):,} from Alpaca")
except Exception as e:
    print(f"    Alpaca failed: {e}")

# ── 3c: SciQ (science Q&A - boost science category) ──
print("  Downloading SciQ...")
try:
    sciq = load_dataset("allenai/sciq", split="train")
    for row in sciq:
        question = (row.get("question") or "").strip()
        answer = (row.get("correct_answer") or "").strip()
        support = (row.get("support") or "").strip()
        if question and answer:
            # Format as instruction with context
            if support:
                out = f"{answer}\n\n{support}" if len(support) > 30 else answer
                new_entries.append(("sciq", {
                    "instruction": question,
                    "input": support[:500] if len(support) > 500 else support,
                    "output": answer
                }))
            else:
                new_entries.append(("sciq", {
                    "instruction": question,
                    "input": "",
                    "output": answer
                }))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='sciq'):,} from SciQ")
except Exception as e:
    print(f"    SciQ failed: {e}")

# ── 3d: GSM8K (math reasoning - boost math category) ──
print("  Downloading GSM8K...")
try:
    gsm = load_dataset("openai/gsm8k", "main", split="train")
    for row in gsm:
        question = (row.get("question") or "").strip()
        answer = (row.get("answer") or "").strip()
        if question and answer:
            new_entries.append(("gsm8k", {
                "instruction": question,
                "input": "",
                "output": answer
            }))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='gsm8k'):,} from GSM8K")
except Exception as e:
    print(f"    GSM8K failed: {e}")

# ── 3e: TruthfulQA (factual accuracy) ──
print("  Downloading TruthfulQA...")
try:
    tqa = load_dataset("truthfulqa/truthful_qa", "generation", split="validation")
    for row in tqa:
        question = (row.get("question") or "").strip()
        best_answer = (row.get("best_answer") or "").strip()
        if question and best_answer:
            new_entries.append(("truthfulqa", {
                "instruction": question,
                "input": "",
                "output": best_answer
            }))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='truthfulqa'):,} from TruthfulQA")
except Exception as e:
    print(f"    TruthfulQA failed: {e}")

# ── 3f: OpenBookQA (science & common sense) ──
print("  Downloading OpenBookQA...")
try:
    obqa = load_dataset("allenai/openbookqa", "main", split="train")
    for row in obqa:
        question = (row.get("question_stem") or "").strip()
        choices = row.get("choices", {})
        answer_key = row.get("answerKey", "")
        if question and choices and answer_key:
            labels = choices.get("label", [])
            texts = choices.get("text", [])
            if answer_key in labels:
                idx = labels.index(answer_key)
                answer = texts[idx]
                # Format nicely
                choices_text = "\n".join(f"{l}) {t}" for l, t in zip(labels, texts))
                new_entries.append(("openbookqa", {
                    "instruction": question,
                    "input": choices_text,
                    "output": f"The answer is {answer_key}) {answer}."
                }))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='openbookqa'):,} from OpenBookQA")
except Exception as e:
    print(f"    OpenBookQA failed: {e}")

# ── 3g: PIQA (physical/common sense reasoning) ──
print("  Downloading PIQA...")
try:
    piqa = load_dataset("ybisk/piqa", split="train", trust_remote_code=True)
    for row in piqa:
        goal = (row.get("goal") or "").strip()
        sol1 = (row.get("sol1") or "").strip()
        sol2 = (row.get("sol2") or "").strip()
        label = row.get("label", -1)
        if goal and sol1 and sol2 and label in [0, 1]:
            answer = sol1 if label == 0 else sol2
            new_entries.append(("piqa", {
                "instruction": f"How would you accomplish the following goal?\n{goal}",
                "input": "",
                "output": answer
            }))
    print(f"    Got {sum(1 for s,_ in new_entries if s=='piqa'):,} from PIQA")
except Exception as e:
    print(f"    PIQA failed: {e}")

total_downloaded = len(new_entries)
print(f"\n  Total downloaded: {total_downloaded:,}")

# ═══════════════════════════════════════════
# STEP 4: Clean new entries
# ═══════════════════════════════════════════
print("\n[4/8] Cleaning new entries...")

# Identity rebranding patterns
REBRAND = [
    (re.compile(r'\bChatGPT\b', re.I), 'LUNA'),
    (re.compile(r'\bGPT-?4o?\b', re.I), 'LUNA'),
    (re.compile(r'\bGPT-?3\.?5?\b', re.I), 'LUNA'),
    (re.compile(r'\bOpenAI\b', re.I), 'Asterizer'),
    (re.compile(r'\bAnthrop(?:ic|ics)\b', re.I), 'Asterizer'),
    (re.compile(r'\bGoogle AI\b', re.I), 'Asterizer'),
    (re.compile(r'\bMeta AI\b', re.I), 'Asterizer'),
    (re.compile(r"(?:I am|I'm) (?:an AI (?:language )?model|a large language model) (?:created|developed|trained) by \w+", re.I),
     'I am LUNA, an AI assistant created by Asterizer'),
]

# Quality filters
NON_ENGLISH_RE = re.compile(r'[\u4e00-\u9fff\u0600-\u06ff\u0900-\u097f\u3040-\u309f\u30a0-\u30ff]')  # CJK, Arabic, Devanagari, Japanese
PLACEHOLDER_RE = re.compile(r'\[insert\b|\[your (?:name|company|topic)\]|<your |_{5,}|XX{3,}', re.I)

cleaned_new = []
rejected = collections.Counter()

for source, entry in new_entries:
    inst = entry["instruction"].strip()
    inp = entry.get("input", "").strip()
    out = entry["output"].strip()
    
    # Skip empty
    if not out or len(out) < 10:
        rejected["too_short_output"] += 1
        continue
    if not inst and not inp:
        rejected["no_instruction_or_input"] += 1
        continue
    
    # Skip non-English
    combined = inst + " " + inp + " " + out
    if NON_ENGLISH_RE.search(combined):
        rejected["non_english"] += 1
        continue
    
    # Skip placeholders
    if PLACEHOLDER_RE.search(combined):
        rejected["placeholder"] += 1
        continue
    
    # Skip if mostly non-ASCII
    ascii_ratio = sum(1 for c in combined if ord(c) < 128) / max(len(combined), 1)
    if ascii_ratio < 0.85:
        rejected["low_ascii"] += 1
        continue
    
    # Rebrand AI identities
    for pattern, replacement in REBRAND:
        inst = pattern.sub(replacement, inst)
        out = pattern.sub(replacement, out)
        if inp:
            inp = pattern.sub(replacement, inp)
    
    # Fix double spaces and basic grammar
    for field_val in [inst, inp, out]:
        field_val = re.sub(r'  +', ' ', field_val)
    
    # Token length check
    entry_clean = {"instruction": inst, "input": inp, "output": out}
    tlen = count_tokens(entry_clean)
    if tlen > MAX_TOKENS:
        rejected["over_token_limit"] += 1
        continue
    if tlen < 15:
        rejected["too_few_tokens"] += 1
        continue
    
    # Dedup against existing
    h = entry_hash(entry_clean)
    if h in existing_hashes:
        rejected["duplicate_existing"] += 1
        continue
    
    existing_hashes.add(h)
    cleaned_new.append((source, entry_clean))

print(f"  Cleaned new entries: {len(cleaned_new):,}")
print(f"  Rejected breakdown:")
for reason, count in rejected.most_common():
    print(f"    {reason}: {count:,}")

# ═══════════════════════════════════════════
# STEP 5: Create input+output only entries (no instruction)
# ═══════════════════════════════════════════
print("\n[5/8] Creating input+output format entries...")

# Convert some entries to input+output only format (reading comprehension style)
input_output_entries = []
# Use some Dolly/Alpaca entries that have context and rephrase
for source, entry in cleaned_new:
    if entry["input"] and len(entry["input"]) > 50 and random.random() < 0.15:
        # Merge instruction into input as a reading prompt
        combined_input = entry["input"]
        if entry["instruction"]:
            combined_input = entry["input"] + "\n\nQuestion: " + entry["instruction"]
        io_entry = {
            "instruction": "",
            "input": combined_input,
            "output": entry["output"]
        }
        tlen = count_tokens(io_entry)
        if tlen <= MAX_TOKENS:
            h = entry_hash(io_entry)
            if h not in existing_hashes:
                existing_hashes.add(h)
                input_output_entries.append(io_entry)

print(f"  Created {len(input_output_entries):,} input+output only entries")

# ═══════════════════════════════════════════
# STEP 6: Diversify instruction phrasing
# ═══════════════════════════════════════════
print("\n[6/8] Diversifying instruction phrasing...")

# For new entries, ensure no instruction prefix repeats >100 times
new_prefix_count = collections.Counter()
diversified = []
for source, entry in cleaned_new:
    prefix = entry["instruction"][:50].lower().strip()
    new_prefix_count[prefix] += 1
    if new_prefix_count[prefix] <= 100:
        diversified.append(entry)

dropped_new_repetition = len(cleaned_new) - len(diversified)
print(f"  Dropped {dropped_new_repetition:,} new entries with repeated prefixes")
cleaned_new_entries = diversified

# ═══════════════════════════════════════════
# STEP 7: Merge everything
# ═══════════════════════════════════════════
print("\n[7/8] Merging all data...")

all_data = existing + cleaned_new_entries + input_output_entries

# Final global dedup
seen = set()
final = []
dup_count = 0
for entry in all_data:
    h = entry_hash(entry)
    if h not in seen:
        seen.add(h)
        final.append(entry)
    else:
        dup_count += 1

print(f"  Final dedup removed: {dup_count:,}")
print(f"  Total final entries: {len(final):,}")

# ═══════════════════════════════════════════
# STEP 8: Token count, split, save, report
# ═══════════════════════════════════════════
print("\n[8/8] Counting tokens and saving...")

total_tokens = 0
for entry in final:
    total_tokens += count_tokens(entry)

avg_tokens = total_tokens / len(final)

# Shuffle and split
random.shuffle(final)
val_size = max(3000, int(len(final) * 0.02))
train = final[val_size:]
val = final[:val_size]

# Save
for path, subset in [
    (OUT_DIR / "all_sft_clean.json", final),
    (OUT_DIR / "train.json", train),
    (OUT_DIR / "val.json", val),
]:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(subset, f, indent=2, ensure_ascii=False)

# Category analysis of final dataset
cats = collections.Counter()
for d in final:
    inst = (d.get("instruction","") + " " + d.get("input","")).lower()
    if any(w in inst for w in ['code','program','function','python','javascript','html','css','sql','api','debug']):
        cats['coding'] += 1
    elif any(w in inst for w in ['math','calculat','equation','algebra','geometry','statistic','probability','solve']):
        cats['math'] += 1
    elif any(w in inst for w in ['science','physics','chemistry','biology','atom','molecule','cell','gravity','evolution','experiment']):
        cats['science'] += 1
    elif any(w in inst for w in ['history','war','century','ancient','empire','civilization','revolution']):
        cats['history'] += 1
    elif any(w in inst for w in ['geography','country','continent','ocean','mountain','river','capital']):
        cats['geography'] += 1
    elif any(w in inst for w in ['health','medicine','disease','symptom','treatment','doctor','nutrition']):
        cats['health'] += 1
    elif any(w in inst for w in ['econom','finance','market','invest','business','trade','gdp','inflation']):
        cats['economics'] += 1
    elif any(w in inst for w in ['write','story','poem','essay','creative','novel','fiction']):
        cats['creative_writing'] += 1
    elif any(w in inst for w in ['explain','what is','define','describe','meaning of']):
        cats['explanation'] += 1
    elif any(w in inst for w in ['summariz','summary','brief','tldr','condense']):
        cats['summarization'] += 1
    elif any(w in inst for w in ['translat']):
        cats['translation'] += 1
    elif any(w in inst for w in ['logic','reason','puzzle','riddle','brain','think step','goal']):
        cats['reasoning'] += 1
    elif any(w in inst for w in ['ethic','moral','philosoph','right wrong','dilemma']):
        cats['philosophy'] += 1
    else:
        cats['general'] += 1

# Format distribution
fmt_inst_out = sum(1 for d in final if d.get("instruction","").strip() and not d.get("input","").strip())
fmt_inst_inp_out = sum(1 for d in final if d.get("instruction","").strip() and d.get("input","").strip())
fmt_inp_out = sum(1 for d in final if not d.get("instruction","").strip() and d.get("input","").strip())

# Source distribution
source_info = {
    "existing (cleaned)": len(existing),
    "dolly": sum(1 for s,_ in cleaned_new if s=='dolly' if _ in cleaned_new_entries),
    "new entries": len(cleaned_new_entries),
    "input+output format": len(input_output_entries),
}

report = f"""LUNA SFT Augmentation & Purification Report
{'='*55}

BEFORE:
  Entries:          304,475
  Tokens:        86,633,618

CLEANING:
  Repetitive prefix cap:   -{removed_for_repetition:,}
  New datasets downloaded:  {total_downloaded:,}
  New entries after clean:  {len(cleaned_new_entries):,}
  Input+output entries:     {len(input_output_entries):,}
  Final dedup removed:      {dup_count:,}

AFTER:
  Total entries:   {len(final):>10,}
  Total tokens:    {total_tokens:>10,}
  Avg tokens/entry:{avg_tokens:>10.1f}
  Train set:       {len(train):>10,}
  Val set:         {len(val):>10,}

FORMAT DISTRIBUTION:
  instruction + output:         {fmt_inst_out:>8,} ({fmt_inst_out/len(final)*100:.1f}%)
  instruction + input + output: {fmt_inst_inp_out:>8,} ({fmt_inst_inp_out/len(final)*100:.1f}%)
  input + output only:          {fmt_inp_out:>8,} ({fmt_inp_out/len(final)*100:.1f}%)

CATEGORY DISTRIBUTION:
"""
for cat, count in cats.most_common():
    report += f"  {cat:20s}: {count:>8,} ({count/len(final)*100:.1f}%)\n"

report += f"""
Files saved:
  {OUT_DIR / 'all_sft_clean.json'}
  {OUT_DIR / 'train.json'}
  {OUT_DIR / 'val.json'}
"""

report_path = OUT_DIR / "AUGMENTATION_REPORT.txt"
with open(report_path, "w", encoding="utf-8") as f:
    f.write(report)

print(report)
print("Done!")