# Updated for ShareGPT
import os
import pandas as pd
import numpy as np
from transformers import AutoTokenizer
import matplotlib.pyplot as plt # Optional, for distribution logic

# ================= CONFIGURATION =================
# Path to your unified parquet file (created by PMPF3)
DATASET_FILE = os.path.join(os.getcwd(), "dataset_cache", "unified_dataset.parquet")
# Path to your model (for the tokenizer)
MODEL_PATH = os.getcwd() 
# =================================================

# =================================================

# Mistral Tekken Templates for exact length scanning
MT_SYSTEM = "[SYSTEM_PROMPT]{sys}[/SYSTEM_PROMPT]"
MT_USER = "[INST] {q} [/INST]"
MT_ASSISTANT = "{a}</s>"

def format_sharegpt_to_text(conversations):
    """
    Mimics how the model sees the data during training.
    Reconstructs the full sequence text from the structured ShareGPT list.
    """
    if not isinstance(conversations, (list, np.ndarray)):
        return ""
        
    text = "<s>"
    for turn in conversations:
        if not isinstance(turn, dict):
            continue
        role = turn.get('from', '')
        content = turn.get('value', '')
        
        if role == 'system':
            text += MT_SYSTEM.format(sys=content)
        elif role == 'human':
            text += MT_USER.format(q=content)
        elif role == 'gpt':
            text += MT_ASSISTANT.format(a=content)
            
    return text

def main():
    print("--- DATASET CONTEXT SCANNER ---")
    
    # 1. Load Tokenizer
    print(f"Loading Tokenizer from {MODEL_PATH}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
    except Exception as e:
        print(f"Error loading tokenizer: {e}")
        print("Cannot calculate exact tokens without tokenizer.")
        return

    # 2. Load Dataset
    if not os.path.exists(DATASET_FILE):
        print(f"Error: {DATASET_FILE} not found.")
        print("Run the Unify script (PMPF3) first.")
        return

    print(f"Loading dataset: {DATASET_FILE}")
    df = pd.read_parquet(DATASET_FILE)
    
    print(f"Scanning {len(df)} entries...")
    
    # 3. Scan
    lengths = []
    char_counts = []
    
    max_tokens = 0
    max_chars = 0
    longest_text = ""
    
    # Iterate and measure
    for index, row in df.iterrows():
        # Retrieve the structured conversations array
        if 'conversations' not in row or row['conversations'] is None:
            continue
            
        conversations = row['conversations']
        
        # Reconstruct the string sequence as tokenized on GPU
        text = format_sharegpt_to_text(conversations)
        if not text:
            continue
            
        # Count characters
        c_len = len(text)
        char_counts.append(c_len)
        
        # Count tokens (Exact)
        # We use len(encode) to get the real number the GPU sees
        ids = tokenizer.encode(text, add_special_tokens=True)
        t_len = len(ids)
        lengths.append(t_len)
        
        if t_len > max_tokens:
            max_tokens = t_len
            max_chars = c_len
            longest_text = text

    # 4. Statistics
    lengths = np.array(lengths)
    p95 = np.percentile(lengths, 95)
    p99 = np.percentile(lengths, 99)
    
    print("\n" + "="*40)
    print("       SCAN RESULTS       ")
    print("="*40)
    print(f"Total Entries:      {len(df)}")
    print(f"Max Characters:     {max_chars}")
    print(f"Max Tokens (Exact): {max_tokens}")
    print("-" * 40)
    print(f"Average Tokens:     {int(np.mean(lengths))}")
    print(f"95th Percentile:    {int(p95)} tokens (Covers 95% of data)")
    print(f"99th Percentile:    {int(p99)} tokens (Covers 99% of data)")
    print("="*40)
    
    # 5. Recommendation
    print("\n--- RECOMMENDATION ---")
    
    # Round up to nearest multiple of 64 for GPU efficiency
    recommended_ctx = int(max_tokens)
    remainder = recommended_ctx % 64
    if remainder != 0:
        recommended_ctx += (64 - remainder)
        
    # Add a tiny safety buffer
    recommended_ctx += 64 
    
    print(f"Set MAX_SEQ_LENGTH = {recommended_ctx}")
    
    if recommended_ctx > 2048:
        print("\n⚠️  WARNING: This is high for an 8GB card.")
        print(f"Consider truncating to {int(p99)} (99th percentile) to save VRAM.")
    else:
        print(f"✅ Safe for 8GB VRAM (assuming LoRA r=16).")

    print("\n--- LONGEST ENTRY PREVIEW ---")
    print(longest_text[:500] + "..." + longest_text[-100:])

if __name__ == "__main__":
    main()