#!/usr/bin/env python3 -u
"""
gen_dataset_figure.py — Generate a 2-panel figure from dataset_stats.csv:
  Left:  Stacked bar chart of sentence counts by script (ar/az/mi)
  Right: Bar chart of median sentence length by script
"""
import csv, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

plt.rcParams.update({
    'font.size': 14, 'axes.titlesize': 16, 'axes.labelsize': 15,
    'xtick.labelsize': 13, 'ytick.labelsize': 13, 'legend.fontsize': 12,
})

CSV_PATH = "/root/oiq_cc_tokenizer/results/dataset_stats.csv"
FIG_DIR = "/root/oiq_cc/figures"
PLOTS_DIR = "/root/oiq_cc_tokenizer/results/plots"

datasets = []
with open(CSV_PATH) as f:
    for r in csv.DictReader(f):
        datasets.append(r)

names = [r["dataset"] for r in datasets]
ar_counts = [int(r["ar_count"]) for r in datasets]
az_counts = [int(r["az_count"]) for r in datasets]
mi_counts = [int(r["mi_count"]) for r in datasets]

ar_medians = [float(r["ar_median"]) if int(r["ar_count"]) > 0 else 0 for r in datasets]
az_medians = [float(r["az_median"]) if int(r["az_count"]) > 0 else 0 for r in datasets]
mi_medians = [float(r["mi_median"]) if int(r["mi_count"]) > 0 else 0 for r in datasets]

ar_means = [float(r["ar_mean"]) if int(r["ar_count"]) > 0 else 0 for r in datasets]
az_means = [float(r["az_mean"]) if int(r["az_count"]) > 0 else 0 for r in datasets]
mi_means = [float(r["mi_mean"]) if int(r["mi_count"]) > 0 else 0 for r in datasets]

COLOR_AR = "#0072B2"
COLOR_AZ = "#D55E00"
COLOR_MI = "#009E73"

fig, axes = plt.subplots(1, 2, figsize=(16, 7))

# Panel 1: Stacked bar chart of sentence counts
ax = axes[0]
x = np.arange(len(names))
width = 0.55

bars_ar = ax.bar(x, ar_counts, width, color=COLOR_AR, edgecolor='black', linewidth=0.5, label='Arabic')
bars_az = ax.bar(x, az_counts, width, bottom=ar_counts, color=COLOR_AZ, edgecolor='black', linewidth=0.5, label='Arabizi')
bottom_mi = [a + b for a, b in zip(ar_counts, az_counts)]
bars_mi = ax.bar(x, mi_counts, width, bottom=bottom_mi, color=COLOR_MI, edgecolor='black', linewidth=0.5, label='Mixed')

ax.set_xticks(x)
ax.set_xticklabels(names, fontsize=11, rotation=20, ha='right')
ax.set_ylabel("Sentence Count", fontweight='bold', fontsize=14)
ax.set_title("Script Distribution by Dataset", fontweight='bold', fontsize=15)
ax.set_yscale('log')
ax.legend(fontsize=12, loc='upper left')
ax.grid(axis='y', alpha=0.3, which='both')

for i, (a, z, m) in enumerate(zip(ar_counts, az_counts, mi_counts)):
    total = a + z + m
    if total > 0:
        ax.text(i, total * 1.15, f"{total:,}", ha='center', va='bottom', fontsize=10, fontweight='bold')

# Panel 2: Grouped bar chart of median sentence length
ax = axes[1]
width = 0.25
x = np.arange(len(names))

bars1 = ax.bar(x - width, ar_medians, width, color=COLOR_AR, edgecolor='black', linewidth=0.5, label='Arabic')
bars2 = ax.bar(x, az_medians, width, color=COLOR_AZ, edgecolor='black', linewidth=0.5, label='Arabizi')
bars3 = ax.bar(x + width, mi_medians, width, color=COLOR_MI, edgecolor='black', linewidth=0.5, label='Mixed')

ax.set_xticks(x)
ax.set_xticklabels(names, fontsize=11, rotation=20, ha='right')
ax.set_ylabel("Median Sentence Length (chars)", fontweight='bold', fontsize=14)
ax.set_title("Median Sentence Length by Script", fontweight='bold', fontsize=15)
ax.legend(fontsize=12)
ax.set_yscale('log')
ax.grid(axis='y', alpha=0.3, which='both')

for bars, vals in [(bars1, ar_medians), (bars2, az_medians), (bars3, mi_medians)]:
    for bar, val in zip(bars, vals):
        if val > 0:
            ax.text(bar.get_x() + bar.get_width()/2, val * 1.15,
                    f"{val:.0f}", ha='center', va='bottom', fontsize=9, fontweight='bold')

plt.tight_layout()
out_path = os.path.join(FIG_DIR, "dataset_comparison.png")
plt.savefig(out_path, dpi=300, bbox_inches='tight')
plt.close()
print(f"Saved: {out_path}")

import shutil
shutil.copy(out_path, os.path.join(PLOTS_DIR, "dataset_comparison.png"))
print(f"Saved: {os.path.join(PLOTS_DIR, 'dataset_comparison.png')}")