import argparse
import json
import os
import glob
import pickle
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoTokenizer
from tqdm import tqdm

# --- Global variables for workers ---
_PICKLE_DATA = None
_TOKENIZER = None

def _load_pickle(path: str):
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Pickle not found: {path}")
    with open(path, "rb") as f:
        return pickle.load(f)

def _load_json_heads(path: str):
    if not os.path.isfile(path):
        raise FileNotFoundError(f"JSON not found: {path}")
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if not isinstance(data, list):
        raise ValueError(f"JSON content must be a list of head names: {path}")
    head_names = []
    for item in data:
        if isinstance(item, list) and item:
            head_names.append(str(item[0]))
        else:
            head_names.append(str(item))
    return head_names

def _as_bool_list(values, name: str, expected_len: int):
    if len(values) != expected_len:
        raise ValueError(f"{name} length {len(values)} != token length {expected_len}")
    return [bool(v) for v in values]

# --- Worker Initializer ---
def init_worker(pickle_path, model_path):
    """
    Called once per process to load heavy resources.
    """
    global _PICKLE_DATA, _TOKENIZER
    
    # Load Pickle
    # print(f"[Worker {os.getpid()}] Loading pickle...")
    _PICKLE_DATA = _load_pickle(pickle_path)
    
    # Load Tokenizer
    # print(f"[Worker {os.getpid()}] Loading tokenizer...")
    _TOKENIZER = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    if _TOKENIZER.pad_token_id is None:
        _TOKENIZER.pad_token = _TOKENIZER.eos_token
        _TOKENIZER.pad_token_id = _TOKENIZER.eos_token_id
        _TOKENIZER.padding_side = "right"

def visualize_heads_task(json_path, output_dir, prompt_index):
    """
    The actual task run by workers.
    """
    global _PICKLE_DATA, _TOKENIZER
    
    try:
        filename = os.path.basename(json_path)
        name_no_ext = os.path.splitext(filename)[0]
        
        # Load heads for this specific task
        head_names = _load_json_heads(json_path)
        if not head_names:
            return f"⏩ Skipped (empty): {filename}"

        # Define Output Path
        png_filename = f"{name_no_ext}.png"
        png_path = os.path.join(output_dir, png_filename)
        
        # --- Visualization Logic (Using Globals) ---
        prompts = _PICKLE_DATA.get("prompts", [])
        heads_data = _PICKLE_DATA.get("heads", {})
        
        if not prompts:
            return f"❌ Error {filename}: No prompts in pickle"
            
        if prompt_index < 0 or prompt_index >= len(prompts):
            return f"❌ Error {filename}: Index {prompt_index} out of range"

        valid_heads = [h for h in head_names if h in heads_data]
        if not valid_heads:
            return f"⚠️ Warning {filename}: No valid heads found"

        prompt = prompts[prompt_index]
        token_ids = list(prompt["token_ids"])
        token_len = len(token_ids)
        
        inst_mask = _as_bool_list(prompt["inst_mask"], "inst_mask", token_len)
        instr_mask = _as_bool_list(prompt["instr_mask"], "instr_mask", token_len)
        user_mask = _as_bool_list(prompt["user_mask"], "user_mask", token_len)

        tokens = [_TOKENIZER.decode([tid], skip_special_tokens=False).replace("\n", "\\n") for tid in token_ids]

        attn_rows = []
        for head_name in valid_heads:
            attn = heads_data[head_name]["attn_weight"][prompt_index]
            if len(attn) != token_len:
                raise ValueError(f"Shape mismatch for {head_name}")
            attn_rows.append(attn)
        
        attn_matrix = np.array(attn_rows, dtype=np.float32).T 

        # --- Plotting ---
        num_heads = len(valid_heads)
        text_panel_width = 4.0
        width_per_head = 0.25
        colorbar_pad = 1.5
        
        heatmap_width = max(2.0, num_heads * width_per_head)
        total_width = text_panel_width + heatmap_width + colorbar_pad
        total_height = max(8, token_len * 0.18)

        fig = plt.figure(figsize=(total_width, total_height))
        gs = fig.add_gridspec(1, 2, width_ratios=[text_panel_width, heatmap_width], wspace=0.05)
        
        ax_text = fig.add_subplot(gs[0, 0])
        ax_heat = fig.add_subplot(gs[0, 1])

        ax_text.set_axis_off()
        ax_text.set_xlim(0, 1)
        ax_text.set_ylim(token_len - 0.5, -0.5)
        ax_text.text(0.0, -1.0, "idx  inst  instr  user  token", fontsize=9, fontfamily="monospace", fontweight='bold')
        
        for i, (t, im, inrm, um) in enumerate(zip(tokens, inst_mask, instr_mask, user_mask)):
            bg_color = "white"
            if um: bg_color = "#e6f2ff"
            elif inrm: bg_color = "#fff2e6"
            elif im: bg_color = "#f2ffe6"
            
            ax_text.text(0.0, i, f"{i:03d}  {int(im):4d}  {int(inrm):5d}  {int(um):4d}  {t}", 
                         fontsize=9, fontfamily="monospace", va="center", 
                         bbox=dict(facecolor=bg_color, edgecolor='none', pad=1))

        im = ax_heat.imshow(attn_matrix, aspect="auto", interpolation="nearest", cmap="viridis", vmin=0, vmax=1.0)
        ax_heat.set_yticks(range(token_len))
        ax_heat.set_yticklabels([f"{i:03d}" for i in range(token_len)], fontsize=7)
        
        max_ticks = 40
        step = max(1, num_heads // max_ticks)
        indices = range(0, num_heads, step)
        labels = [valid_heads[i] for i in indices]

        ax_heat.set_xticks(indices)
        ax_heat.set_xticklabels(labels, rotation=90, ha="center", fontsize=8)
        
        ax_heat.set_xlabel(f"Top {num_heads} Heads")
        ax_heat.set_ylabel("Tokens")
        
        cbar = fig.colorbar(im, ax=ax_heat, fraction=0.046, pad=0.04)
        cbar.set_label("Attention Weight")

        fig.suptitle(f"Setup: {name_no_ext}\n(Prompt Index: {prompt_index})", y=0.99, fontsize=12)

        fig.savefig(png_path, bbox_inches="tight", dpi=150)
        plt.close(fig)
        
        return None # Success

    except Exception as e:
        traceback.print_exc()
        return f"❌ Exception in {json_path}: {str(e)}"

def main():
    parser = argparse.ArgumentParser(description="Batch visualize sorted heads (Multiprocessing).")
    parser.add_argument("--input-dir", required=True, help="Directory containing sorted JSON files")
    parser.add_argument("--pickle-path", required=True, help="Path to raw pickle")
    parser.add_argument("--prompt-index", type=int, default=0)
    parser.add_argument("--model-path", default="../../models/Llama-3.1-8B-Instruct")
    parser.add_argument("--workers", type=int, default=int(os.cpu_count()/4) or 4)
    args = parser.parse_args()

    input_abs = os.path.abspath(args.input_dir)
    parent_dir = os.path.dirname(input_abs)
    output_dir = os.path.join(parent_dir, "heads_sorted_visualize")
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"📂 Reading heads from: {input_abs}")
    print(f"💾 Output images to: {output_dir}")
    print(f"🚀 Using {args.workers} workers")

    json_pattern = os.path.join(input_abs, "*.json")
    json_files = sorted(glob.glob(json_pattern))
    
    if not json_files:
        print("❌ No JSON files found.")
        return

    # Using ProcessPoolExecutor to bypass GIL for Matplotlib
    with ProcessPoolExecutor(
        max_workers=args.workers,
        initializer=init_worker,
        initargs=(args.pickle_path, args.model_path)
    ) as executor:
        
        futures = {executor.submit(visualize_heads_task, jf, output_dir, args.prompt_index): jf for jf in json_files}
        
        for future in tqdm(as_completed(futures), total=len(json_files), desc="Rendering"):
            result = future.result()
            if result: # If string returned, it's an error/warning message
                print(result)

    print("\n✅ Batch visualization complete.")

if __name__ == "__main__":
    main()
