import os
import glob
import pickle
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns


# ==========================================
# 1. DATA LOADING & PROCESSING
# ==========================================

def load_data_for_task(base_dir, task_name, model_pattern):
    """
    Loads pickle files matching the model_pattern for a specific task.
    Returns:
        - samples_data: dict of sample-level data (scores, visual tokens)
        - metadata: dict containing 'e2e_latency', 'total_tokens', etc.
    """
    # Construct search pattern. logic handles wildcards if passed in model_pattern
    search_pattern = os.path.join(base_dir, f"{task_name}_analysis_rank*_{model_pattern}*.pkl")
    files = glob.glob(search_pattern)

    if not files:
        print(f"  [Warning] No files found for task '{task_name}' with pattern '*{model_pattern}*'")
        return {}, {}

    merged_samples = {}
    aggregated_metadata = {"e2e_latency": 0, "total_tokens": 0}
    found_metadata = False

    for file_path in files:
        try:
            with open(file_path, "rb") as f:
                data_chunk = pickle.load(f)

                # Separate metric_dict from sample data
                if 'metric_dict' in data_chunk:
                    m_dict = data_chunk.pop('metric_dict')
                    # Sum latencies/tokens if split across multiple files (rare but safe)
                    aggregated_metadata['e2e_latency'] += m_dict.get('e2e_latency', 0)
                    aggregated_metadata['total_tokens'] += m_dict.get('total_tokens', 0)
                    found_metadata = True

                merged_samples.update(data_chunk)
        except Exception as e:
            print(f"  [Error] Failed to load {file_path}: {e}")

    if not found_metadata:
        print(f"  [Info] No 'metric_dict' found in files for {task_name}.")

    return merged_samples, aggregated_metadata


def extract_metrics(task_data_map, task_meta_map):
    """
    Calculates average metrics (Throughput, Avg Visual Tokens) from raw data.
    """
    visual_tokens = []

    # 1. Extract Visual Token Counts
    for doc_id, content in task_data_map.items():
        if "visual_token_num" in content:
            val = content["visual_token_num"]
            # Handle if it's a tensor/array or scalar
            if hasattr(val, 'item'):
                val = val.item()
            visual_tokens.append(float(val))

    avg_vis_tokens = np.mean(visual_tokens) if visual_tokens else 0

    # 2. Calculate Throughput (Samples / Second)
    # Total samples in this dataset
    n_samples = len(task_data_map)

    # Total time (e2e_latency is usually in seconds, ensure your logger units match)
    total_time = task_meta_map.get('e2e_latency', 0)

    if total_time > 0:
        throughput = n_samples / total_time
    else:
        throughput = 0

    return {
        "avg_visual_tokens": avg_vis_tokens,
        "throughput": throughput,
        "n_samples": n_samples
    }


# ==========================================
# 2. PLOTTING FUNCTIONS
# ==========================================

def plot_tradeoff_curves(results, save_dir):
    """
    Plots Accuracy vs. Efficiency (Throughput) and Accuracy vs. Visual Tokens.
    Handles Baseline, Ours, and RoI curves with distinct styles.
    """
    labels = list(results.keys())

    # Calculate Macro Averages (across all tasks for each config)
    avg_accs = []
    avg_throughputs = []
    avg_tokens = []

    for label in labels:
        metrics = results[label]
        # Calculate mean of values; ensure we handle potential empty lists safely
        acc_values = [m['acc'] for m in metrics.values()]
        th_values = [m['throughput'] for m in metrics.values()]
        tok_values = [m['avg_visual_tokens'] for m in metrics.values()]

        avg_accs.append(np.mean(acc_values) if acc_values else 0)
        avg_throughputs.append(np.mean(th_values) if th_values else 0)
        avg_tokens.append(np.mean(tok_values) if tok_values else 0)

    # --- STYLE DEFINITIONS ---
    # Customize colors and markers here
    styles = {
        'Baseline': {'marker': 'o', 'linestyle': '--', 'color': '#7f8c8d'},  # Grey, Dashed, Circle
        'Ours': {'marker': '*', 'linestyle': '-', 'color': '#e74c3c'},  # Red, Solid, Star
        'RoI': {'marker': 's', 'linestyle': '-', 'color': '#3498db'}  # Blue, Solid, Square
    }

    # Helper function to generate both plots without code duplication
    def create_plot(x_data, x_label, file_name, title):
        plt.figure(figsize=(10, 7))
        sns.set_theme(style="whitegrid")  # Optional: makes grid look nice

        # Identify unique prefixes (Baseline, Ours, RoI) from the config keys
        prefixes = sorted(list(set([l.split('\n')[0] for l in labels])))

        # Sort prefixes to ensure Legend order (e.g. Baseline -> Ours -> RoI) if desired
        # You can manually force order: prefixes = ['Baseline', 'Ours', 'RoI']

        for prefix in prefixes:
            # Get style, default to black if unknown prefix
            style = styles.get(prefix, {'marker': 'D', 'linestyle': '-.', 'color': 'black'})

            # Filter data for this group
            group_indices = [i for i, l in enumerate(labels) if l.startswith(prefix)]
            if not group_indices: continue

            g_x = [x_data[i] for i in group_indices]
            g_y = [avg_accs[i] for i in group_indices]
            g_labels = [labels[i] for i in group_indices]

            # Sort points by X-axis (efficiency or tokens) to draw correct line
            sorted_pairs = sorted(zip(g_x, g_y, g_labels), key=lambda x: x[0])
            s_x, s_y, s_l = zip(*sorted_pairs)

            # Plot Line
            plt.plot(s_x, s_y,
                     marker=style['marker'],
                     linestyle=style['linestyle'],
                     color=style['color'],
                     label=prefix,
                     linewidth=2,
                     markersize=10 if style['marker'] != '*' else 14)  # Stars need to be larger

            # Annotate Points (Labels like "Max-256")
            for x, y, lbl in zip(s_x, s_y, s_l):
                clean_lbl = lbl.split('\n')[-1]

                # Dynamic offsets to prevent overlapping labels
                offset = (0, 10)
                if prefix == 'RoI':
                    offset = (0, -15)
                elif prefix == 'Ours':
                    offset = (-10, 10)

                plt.annotate(clean_lbl, (x, y), xytext=offset, textcoords='offset points',
                             ha='center', fontsize=8, fontweight='bold',
                             bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.6, ec="none"))

        plt.title(title, fontsize=14)
        plt.xlabel(x_label, fontweight='bold')
        plt.ylabel('Average Accuracy (%)', fontweight='bold')
        plt.grid(True, linestyle=':', alpha=0.6)
        plt.legend(fontsize=10)
        plt.tight_layout()

        save_path = os.path.join(save_dir, file_name)
        plt.savefig(save_path, dpi=300)
        plt.close()
        print(f"Saved plot: {save_path}")

    # 1. Plot Efficiency
    create_plot(avg_throughputs,
                'Average Throughput (Samples / Second) \n← Slower | Faster →',
                "acc_vs_efficiency.png",
                "Accuracy vs. Efficiency Trade-off")

    # 2. Plot Visual Tokens
    create_plot(avg_tokens,
                'Average Visual Tokens per Image \n← Cheaper | More Detailed →',
                "acc_vs_visual_tokens.png",
                "Accuracy vs. Visual Cost")

# ==========================================
# 3. EXISTING UTILS (Preserved)
# ==========================================
def print_threshold_ratios(task_data_map, threshold):
    print(f"\n{'=' * 20} Threshold Analysis (Score > {threshold}) {'=' * 20}")
    print(f"{'Dataset Name':<20} | {'Total Samples':<15} | {'Above Thr':<10} | {'Ratio':<10}")
    print("-" * 65)
    for task_name, content in task_data_map.items():
        scores = content['scores']
        if len(scores) == 0: continue
        above_mask = scores > threshold
        count_above = np.sum(above_mask)
        ratio = count_above / len(scores)
        print(f"{task_name:<20} | {len(scores):<15} | {count_above:<10} | {ratio:.2%}")
    print("-" * 65 + "\n")


def extract_scores_list(data_dict):
    scores = []
    for doc_id, content in data_dict.items():
        if "high_res_pred_score" in content:
            val = content["high_res_pred_score"]
            if isinstance(val, np.ndarray): val = val.flatten()[0]
            scores.append(float(val))
    return np.array(scores)


def plot_cost_benefit_combo(results, task_name, save_dir):
    """
    Better Idea: Combines Relative Efficiency (Bars) with Accuracy Gain (Points/Lines)
    on a dual-axis chart.
    """
    # 1. Parse Data (Similar to above, but also gets Accuracy)
    data_map = {}
    for label, metrics in results.items():
        method, resolution = label.split('\n')
        if resolution not in data_map: data_map[resolution] = {}

        task_data = metrics.get(task_name, {'throughput': 0, 'acc': 0})
        data_map[resolution][method] = {
            'speed': task_data['throughput'],
            'acc': task_data['acc']
        }

    resolutions = sorted(data_map.keys(), key=lambda x: int(x.split('-')[-1]))
    methods = ['Baseline', 'Ours', 'RoI']
    colors = {'Baseline': '#95a5a6', 'Ours': '#e74c3c', 'RoI': '#3498db'}

    x = np.arange(len(resolutions))
    width = 0.25

    fig, ax1 = plt.subplots(figsize=(11, 7))
    ax2 = ax1.twinx()  # Create secondary y-axis for Accuracy Gain

    # Loop for Bars (Efficiency) AND Points (Accuracy)
    for i, method in enumerate(methods):
        eff_ratios = []
        acc_gains = []

        for res in resolutions:
            base_speed = data_map[res].get('Baseline', {}).get('speed', 1.0)
            base_acc = data_map[res].get('Baseline', {}).get('acc', 0.0)

            curr = data_map[res].get(method, {'speed': 0, 'acc': 0})

            # Efficiency Ratio
            ratio = (curr['speed'] / base_speed) * 100 if base_speed > 0 else 0
            eff_ratios.append(ratio)

            # Accuracy Delta (Gain over Baseline of SAME resolution)
            delta = curr['acc'] - base_acc
            acc_gains.append(delta)

        offset = (i - 1) * width

        # PLOT 1: Efficiency Bars (Left Axis)
        ax1.bar(x + offset, eff_ratios, width, color=colors[method], alpha=0.3, label=f"{method} (Speed)")

        # PLOT 2: Accuracy Gain Markers (Right Axis)
        # We plot these slightly offset to align with bars, or connected with lines
        if method == 'Baseline':
            # Plot baseline as flat line at 0
            ax2.plot(x + offset, acc_gains, color=colors[method], marker='_', markersize=20, linestyle='None')
        else:
            # Plot gains as distinct markers
            ax2.plot(x + offset, acc_gains, color=colors[method], marker='D', markersize=8,
                     linestyle='-', linewidth=2, label=f"{method} (Acc Gain)")

            # Annotate specific gains
            for j, gain in enumerate(acc_gains):
                ax2.annotate(f"+{gain:.1f}", (x[j] + offset, gain), xytext=(0, 8),
                             textcoords='offset points', ha='center', color=colors[method], fontweight='bold')

    # Formatting
    ax1.set_ylabel("Relative Efficiency (%)", fontsize=12, fontweight='bold')
    ax1.set_ylim(0, 120)
    ax1.set_xlabel("Resolution", fontsize=12)
    ax1.set_xticks(x)
    ax1.set_xticklabels(resolutions)

    ax2.set_ylabel("Accuracy Gain over Baseline (Points)", fontsize=12, fontweight='bold')
    # Set symmetric limits or auto limits for acc gain
    ax2.axhline(0, color='gray', linestyle=':', linewidth=1)

    plt.title(f"Cost-Benefit Analysis: {task_name}\n(Bars = Speed %, Points = Acc Gain)", fontsize=14)

    # Custom Legend
    lines, labels = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines + lines2, labels + labels2, loc='upper left')

    plt.tight_layout()
    save_path = os.path.join(save_dir, f"{task_name}_cost_benefit.png")
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"Saved Cost-Benefit plot: {save_path}")

# ==========================================
# 4. MAIN EXECUTION
# ==========================================

if __name__ == "__main__":
    BASE_DIR = "../archive_data/"
    SAVE_DIR = "./analysis_plots/"
    os.makedirs(SAVE_DIR, exist_ok=True)

    # --- CONFIGURATION: Define Models and Manual Accuracies ---
    # Keys: Display Label (First line is Group, second is variant)
    # Values: 'pattern': filename part to match, 'accs': dict of {task: accuracy}
    TARGET_TASKS = [
        "docvqa_val",
        "infovqa_val",
        "textvqa_val",
        "chartqa",  # Added from table
        "ocrbench"  # Added from table
        # "vstar_bench",
        # "mmerealworld_lite",
    ]

    # --- CONFIGURATION: Define Models and Manual Accuracies ---
    # Patterns are derived from: Visual_Tokens * 28 * 28
    # 128 tokens = 100352 pixels
    # 256 tokens = 200704 pixels
    # 576 tokens = 451584 pixels
    # 1024 tokens = 802816 pixels
    # 2048 tokens = 1605632 pixels

    MODEL_CONFIGS = {
        # --- BASELINE CONFIGURATIONS ---
        # 'Baseline\nMax-256': {
        #     'pattern': 'Qwen2.5-VL-7B-Instruct_min100352_max200704',
        #     'accs': {
        #         'docvqa_val': 81.2,
        #         'chartqa': 76.8,
        #         'ocrbench': 77.8,
        #         'infovqa_val': 54.0,
        #         'textvqa_val': 75.5,
        #         'vstar_bench': 53.4,
        #         'mmerealworld_lite': 27.6
        #     }
        # },
        'Baseline\nMax-576': {
            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max451584',
            'accs': {
                'docvqa_val': 92.0,
                'chartqa': 83.0,
                'ocrbench': 82.8,
                'infovqa_val': 70.1,
                'textvqa_val': 81.1,
                'vstar_bench': 64.4,
                'mmerealworld_lite': 35.4
            }
        },
        'Baseline\nMax-1024': {
            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max802816',
            'accs': {
                'docvqa_val': 93.9,
                'chartqa': 84.6,
                'ocrbench': 84.4,
                'infovqa_val': 76.5,
                'textvqa_val': 82.8,
                'vstar_bench': 70.2,
                'mmerealworld_lite': 37.5
            }
        },
        'Baseline\nMax-2048': {
            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max1605632',
            'accs': {
                'docvqa_val': 94.7,
                'chartqa': 84.5,
                'ocrbench': 84.8,
                'infovqa_val': 80.4,
                'textvqa_val': 82.9,
                'vstar_bench': 72.3,
                'mmerealworld_lite': 39.3
            }
        },
        #
        # # # # --- YOUR "OURS" CONFIGURATIONS (Update these patterns as needed) ---
        # 'Ours\nMax-256': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min100352_max200704',
        #     'accs': {
        #         'docvqa_val': 89.1,
        #         'chartqa': 82.5,
        #         'ocrbench': 83.8,
        #         'infovqa_val': 69.6,
        #         'textvqa_val': 81.8,
        #         'vstar_bench': 67.5,
        #         'mmerealworld_lite': 37.3
        #     }
        # },
        # 'Ours\nMax-576': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max451584',
        #     'accs': {
        #         'docvqa_val': 93.9,
        #         'chartqa': 84.2,
        #         'ocrbench': 84.8,
        #         'infovqa_val': 78.5,
        #         'textvqa_val': 82.8,
        #         'vstar_bench': 80.1,
        #         'mmerealworld_lite': 43.4
        #     }
        # },
        # 'Ours\nMax-1024': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max802816',
        #     'accs': {
        #         'docvqa_val': 94.6,
        #         'chartqa': 85.4,
        #         'ocrbench': 84.8,
        #         'infovqa_val': 81.5,
        #         'textvqa_val': 83.7,
        #         'vstar_bench': 83.3,
        #         'mmerealworld_lite': 46.5
        #     }
        # },
        # 'Ours\nMax-1536': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max1204224',
        #     'accs': {
        #         'docvqa_val': 94.8,
        #         'chartqa': 85.4,
        #         'ocrbench': 85.4,
        #         'infovqa_val': 82.4,
        #         'textvqa_val': 83.6,
        #         'vstar_bench': 88.0,
        #         'mmerealworld_lite': 47.2
        #     }
        # },
        ###--- "OURS-RoI" CONFIGURATIONS ---
        'RoI\nMax-256': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min100352_max200704',
            'accs': {
                'docvqa_val': 89.2, 'chartqa': 83.0, 'ocrbench': 83.6, 'infovqa_val': 69.9,
                'textvqa_val': 81.4, 'vstar_bench': 67.5, 'mmerealworld_lite': 37.6
            }
        },
        'RoI\nMax-576': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max451584',
            'accs': {
                'docvqa_val': 94.1, 'chartqa': 85.8, 'ocrbench': 84.9, 'infovqa_val': 79.6,
                'textvqa_val': 82.9, 'vstar_bench': 80.6, 'mmerealworld_lite': 44.6
            }
        },
        'RoI\nMax-1024': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max802816',
            'accs': {
                'docvqa_val': 94.7, 'chartqa': 86.3, 'ocrbench': 84.8, 'infovqa_val': 82.0,
                'textvqa_val': 83.9, 'vstar_bench': 82.2, 'mmerealworld_lite': 46.7
            }
        },
        # 'RoI\nMax-1536': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max1204224',
        #     'accs': {
        #         'docvqa_val': 94.9, 'chartqa': 85.8, 'ocrbench': 84.8, 'infovqa_val': 83.2,
        #         'textvqa_val': 83.9, 'vstar_bench': 88.0, 'mmerealworld_lite': 47.0
        #     }
        # },

        # # # --- YOUR "OURS" CONFIGURATIONS (Update these patterns as needed) ---
        'Ours\nMax-1024': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min200704_max802816',
            'accs': {
                'docvqa_val': 94.8,
                'chartqa': 86.4,
                'ocrbench': 85.4,
                'infovqa_val': 81.2,
                'textvqa_val': 84.1,
                'vstar_bench': 83.8,
                'mmerealworld_lite': 47.3
            }
        },
        'Ours\nMax-576': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min200704_max451584',
            'accs': {
                'docvqa_val': 94.0,
                'chartqa': 85.4,
                'ocrbench': 85.4,
                'infovqa_val': 79.4,
                'textvqa_val': 83.5,
                'vstar_bench': 79.6,
                'mmerealworld_lite': 45.7
            }
        },
        'Ours\nMax-384': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min100352_max301056',
            'accs': {
                'docvqa_val': 93.4,
                'chartqa': 84.2,
                'ocrbench': 85.3,
                'infovqa_val': 75.7,
                'textvqa_val': 82.7,
                'vstar_bench': 73.8,
                'mmerealworld_lite': 41.2
            }
        },
        'Ours\nMax-256': {
            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min100352_max200704',
            'accs': {
                'docvqa_val': 90.4,
                'chartqa': 82.7,
                'ocrbench': 83.5,
                'infovqa_val': 71.7,
                'textvqa_val': 81.8,
                'vstar_bench': 68.6,
                'mmerealworld_lite': 37.9
            }
        },
        # 'Ours\nMax-1536': {
        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-post7k-repos-full2_min200704_max1204224',
        #     'accs': {
        #         'docvqa_val': 95.0,
        #         'chartqa': 86.4,
        #         'ocrbench': 84.8,
        #         'infovqa_val': 82.6,
        #         'textvqa_val': 84.3,
        #         'vstar_bench': 86.9,
        #         'mmerealworld_lite': 48.2
        #     }
        # },

    }
    print("--- Loading Data for Comparison ---")

    # Store all results here: { "Baseline\nMax-256": { "textvqa": {throughput, tokens, acc}, ... } }
    overall_results = {}

    for label, config in MODEL_CONFIGS.items():
        print(f"Processing config: {label.replace(chr(10), ' ')}")  # chr(10) is newline
        overall_results[label] = {}

        # Load data for each task for this specific config
        for task in TARGET_TASKS:
            # 1. Load Raw Data
            raw_data, metadata = load_data_for_task(BASE_DIR, task, config['pattern'])

            if not raw_data:
                # Fill dummy if missing to avoid crash, or skip
                overall_results[label][task] = {'throughput': 0, 'avg_visual_tokens': 0, 'acc': 0}
                continue

            # 2. Extract Metrics (Speed, Tokens)
            metrics = extract_metrics(raw_data, metadata)

            # 3. Add Manual Accuracy
            acc = config['accs'].get(task, 0.0)  # Default 0 if task missing in config

            # 4. Store Scores for Histograms (Optional, if you want detailed score plots per config)
            # scores = extract_scores_list(raw_data)

            # Save combined
            overall_results[label][task] = {
                'throughput': metrics['throughput'],
                'avg_visual_tokens': metrics['avg_visual_tokens'],
                'acc': acc,
                # 'scores': scores # Uncomment if you want to pass this to the histogram plotter
            }

            print(
                f"  > {task}: {metrics['throughput']:.1f} samp/s, {metrics['avg_visual_tokens']:.0f} tokens, Acc: {acc}")

    # --- PLOT COMPARISONS ---
    if overall_results:
        print("\n--- Generating Comparison Plots ---")
        plot_tradeoff_curves(overall_results, SAVE_DIR)

    print("\nDone! Plots saved to:", SAVE_DIR)

    for i in range(len(TARGET_TASKS)):
        plot_cost_benefit_combo(overall_results, TARGET_TASKS[i], SAVE_DIR)