#!/usr/bin/env python3
import argparse
import glob
import json
import math
import os
import sys

import numpy as np
from scipy.ndimage import gaussian_filter, distance_transform_edt
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, Normalize
from matplotlib.cm import ScalarMappable


STROKE_COLORS = ["#0000ff", "#00ffff", "#00ff00", "#ffff00", "#ff0000"]
STROKE_CMAP = LinearSegmentedColormap.from_list("strokes", STROKE_COLORS, N=256)
META_CHARS = set("{}^&<>~|\\")


def to_unit(values, scale, lo, hi):
    values = np.asarray(values, dtype=float)
    if scale == "log":
        numerator = np.log(values) - math.log(lo)
        denominator = math.log(hi) - math.log(lo)
    else:
        numerator = values - lo
        denominator = hi - lo
    return np.clip(numerator / denominator, 0.0, 1.0)


def format_rank(rank):
    rank = int(round(rank))
    if rank >= 1000:
        thousands = rank / 1000.0
        if abs(thousands - round(thousands)) < 0.05:
            return f"{thousands:.0f}k"
        return f"{thousands:.1f}k"
    return str(rank)


def format_length(value):
    return str(int(round(value)))


def is_meta_entry(word):
    return any(c in META_CHARS for c in word) or not any(c.isalpha() for c in word)


def load_frequency(path):
    exact_rank, lower_rank = {}, {}
    count = 0
    with open(path, encoding="utf-8", errors="replace") as f:
        for line in f:
            token = line.strip("\r\n")
            if not token:
                continue
            count += 1
            exact_rank.setdefault(token, count)
            lower_rank.setdefault(token.lower(), count)
    return exact_rank, lower_rank, count


def lookup_rank(word, exact_rank, lower_rank):
    return exact_rank.get(word) or lower_rank.get(word.lower())


def collect_strokes(json_paths, filter_meta=True):
    fewest_strokes = {}
    for path in json_paths:
        with open(path, encoding="utf-8", errors="replace") as f:
            dictionary = json.load(f)
        for outline, word in dictionary.items():
            if filter_meta and is_meta_entry(word):
                continue
            strokes = outline.count("/") + 1
            if word not in fewest_strokes or strokes < fewest_strokes[word]:
                fewest_strokes[word] = strokes
    return fewest_strokes


def build_points(fewest_strokes, exact_rank, lower_rank, rarity_axis, length_mode="chars"):
    lengths, ranks, strokes = [], [], []
    for word, stroke_count in fewest_strokes.items():
        length = len(word) if length_mode == "chars" else len(word.split())
        if length <= 0:
            continue
        rank = lookup_rank(word, exact_rank, lower_rank)
        if rank is None:
            if not rarity_axis["include_unranked"]:
                continue
            rank = rarity_axis["n_ranked"] + 1
        lengths.append(length)
        ranks.append(rank)
        strokes.append(stroke_count)

    lengths = np.asarray(lengths, dtype=float)
    ranks = np.asarray(ranks, dtype=float)
    strokes = np.asarray(strokes, dtype=float)
    rarity = to_unit(ranks, rarity_axis["scale"], 1, rarity_axis["max_rank"]) if ranks.size else ranks
    return lengths, rarity, strokes


def aggregate_grid(x01, y01, strokes, grid, how="mean"):
    col = np.clip((x01 * grid).astype(int), 0, grid - 1)
    row = np.clip((y01 * grid).astype(int), 0, grid - 1)
    cell = col * grid + row
    count = np.bincount(cell, minlength=grid * grid).astype(float)

    if how == "mean":
        total = np.bincount(cell, weights=strokes, minlength=grid * grid)
        stat = np.where(count > 0, total / np.maximum(count, 1), np.nan)
    elif how == "min":
        stat = np.full(grid * grid, np.inf)
        np.minimum.at(stat, cell, strokes)
        stat[np.isinf(stat)] = np.nan
    else:
        stat = np.full(grid * grid, np.nan)
        order = np.argsort(cell, kind="stable")
        cells_sorted, strokes_sorted = cell[order], strokes[order]
        splits = np.flatnonzero(np.diff(cells_sorted)) + 1
        starts = cells_sorted[np.concatenate(([0], splits))]
        for group, key in zip(np.split(strokes_sorted, splits), starts):
            stat[key] = np.median(group)

    return stat.reshape(grid, grid), count.reshape(grid, grid)


def smooth_field(stat_grid, count_grid, sigma):
    stat = np.nan_to_num(stat_grid, nan=0.0)
    if sigma > 0:
        weighted = gaussian_filter(stat * count_grid, sigma, mode="nearest")
        weight = gaussian_filter(count_grid, sigma, mode="nearest")
    else:
        weighted = stat * count_grid
        weight = count_grid.astype(float)
    value = np.where(weight > 1e-9, weighted / np.maximum(weight, 1e-9), np.nan)
    return value, weight


def nearest_fill(grid):
    nearest = distance_transform_edt(np.isnan(grid), return_distances=False, return_indices=True)
    return grid[tuple(nearest)]


def style_dark(fig, ax):
    fig.patch.set_facecolor("black")
    ax.set_facecolor("black")
    for spine in ax.spines.values():
        spine.set_color("#666666")
    ax.tick_params(colors="#dddddd")
    ax.xaxis.label.set_color("#ffffff")
    ax.yaxis.label.set_color("#ffffff")
    ax.title.set_color("#ffffff")


def add_colorbar(fig, ax, max_strokes):
    mappable = ScalarMappable(norm=Normalize(vmin=1, vmax=max_strokes), cmap=STROKE_CMAP)
    mappable.set_array([])
    colorbar = fig.colorbar(mappable, ax=ax, pad=0.02, fraction=0.046)
    ticks = list(range(1, max_strokes + 1))
    colorbar.set_ticks(ticks)
    labels = [str(t) for t in ticks]
    labels[-1] = f"{max_strokes}+"
    colorbar.set_ticklabels(labels)
    colorbar.set_label("Strokes per word  (min per word, aggregated per cell)", color="#ffffff")
    colorbar.ax.yaxis.set_tick_params(color="#dddddd")
    plt.setp(plt.getp(colorbar.ax.axes, "yticklabels"), color="#dddddd")


def set_length_ticks(ax, length_axis):
    scale, lo, hi = length_axis["scale"], length_axis["lo"], length_axis["hi"]
    if scale == "log":
        powers = []
        power = 1
        while power <= hi:
            if power >= lo:
                powers.append(power)
            power *= 2
        values = sorted(set([int(round(lo))] + powers + [int(round(hi))]))
    else:
        low, high = int(math.floor(lo)), int(math.ceil(hi))
        step = max(1, (high - low) // 12)
        values = list(range(low, high + 1, step))
    ax.set_xticks([float(to_unit(v, scale, lo, hi)) for v in values])
    ax.set_xticklabels([format_length(v) for v in values])


def set_rarity_ticks(ax, rarity_axis):
    scale = rarity_axis["scale"]
    n_ranked, max_rank = rarity_axis["n_ranked"], rarity_axis["max_rank"]
    ticks, labels = [], []
    if scale == "log":
        rank = 1
        while rank <= n_ranked:
            pos = float(to_unit(rank, scale, 1, max_rank))
            if pos <= 0.97:
                ticks.append(pos)
                labels.append(format_rank(rank))
            rank *= 10
    else:
        for frac in (0.0, 0.2, 0.4, 0.6, 0.8):
            ticks.append(frac)
            labels.append(format_rank(1 + frac * (max_rank - 1)))
    ticks.append(1.0)
    labels.append("unranked" if rarity_axis["include_unranked"] else format_rank(n_ranked))
    ax.set_yticks(ticks)
    ax.set_yticklabels(labels)


def finish_axes(ax, length_axis, rarity_axis, title):
    unit = "characters" if length_axis["mode"] == "chars" else "words"
    ax.set_xlabel(f"Word length ({unit}, {length_axis['scale']} scale)  -  short -> long")
    ax.set_ylabel(f"Word rarity - frequency rank ({rarity_axis['scale']} scale)\n"
                  "more common -> rarer  (bottom -> top)")
    ax.set_title(title, pad=14, fontsize=14)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    set_length_ticks(ax, length_axis)
    set_rarity_ticks(ax, rarity_axis)
    ax.text(0.012, 0.02, "common * short", transform=ax.transAxes,
            color="#bbbbbb", fontsize=8, ha="left", va="bottom", alpha=0.8)
    ax.text(0.988, 0.98, "rare * long", transform=ax.transAxes,
            color="#bbbbbb", fontsize=8, ha="right", va="top", alpha=0.8)


def make_axes(args):
    fig, ax = plt.subplots(figsize=(args.width, args.height), dpi=args.dpi)
    style_dark(fig, ax)
    norm = Normalize(vmin=1, vmax=args.max_strokes)
    return fig, ax, norm


def finalize(fig, ax, length_axis, rarity_axis, args):
    add_colorbar(fig, ax, args.max_strokes)
    finish_axes(ax, length_axis, rarity_axis, args.title)
    fig.tight_layout()
    fig.savefig(args.out, facecolor="black", bbox_inches="tight")
    plt.close(fig)


def render_heatmap(x, y, strokes, length_axis, rarity_axis, args):
    x01 = to_unit(x, length_axis["scale"], length_axis["lo"], length_axis["hi"])
    stat_grid, count_grid = aggregate_grid(x01, y, strokes, args.grid, how=args.aggregate)
    value, density = smooth_field(stat_grid, count_grid, args.smooth)
    fig, ax, norm = make_axes(args)

    if args.shade == "solid":
        filled = nearest_fill(value)
        if args.smooth > 0:
            filled = gaussian_filter(filled, args.smooth, mode="nearest")
        rgb = STROKE_CMAP(norm(np.clip(filled, 1, args.max_strokes)))[..., :3]
    else:
        rgb = STROKE_CMAP(norm(np.clip(value, 1, args.max_strokes)))[..., :3]
        if args.shade == "mask":
            brightness = (density > 1e-9).astype(float)
        else:
            positive = density[density > 1e-9]
            reference = np.percentile(positive, 99) if positive.size else 1.0
            brightness = np.clip(np.log1p(density) / math.log1p(max(reference, 1e-9)), 0, 1) ** args.gamma
        rgb = rgb * brightness[..., None]
        rgb[np.isnan(value)] = 0.0

    ax.imshow(np.transpose(rgb, (1, 0, 2)), origin="lower", extent=[0, 1, 0, 1],
              aspect="auto", interpolation="bilinear")
    finalize(fig, ax, length_axis, rarity_axis, args)


def render_scatter(x, y, strokes, length_axis, rarity_axis, args):
    x01 = to_unit(x, length_axis["scale"], length_axis["lo"], length_axis["hi"])
    jitter = (np.random.default_rng(0).random(x01.size) - 0.5) * 0.02
    fig, ax, norm = make_axes(args)
    ax.scatter(np.clip(x01 + jitter, 0, 1), y, c=np.clip(strokes, 1, args.max_strokes),
               cmap=STROKE_CMAP, norm=norm, s=args.point_size, alpha=args.point_alpha, linewidths=0)
    finalize(fig, ax, length_axis, rarity_axis, args)


def parse_args(argv=None):
    p = argparse.ArgumentParser(
        description="Heat-map of Plover steno stroke count vs word length x rarity.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    io = p.add_argument_group("input / output")
    io.add_argument("--dir", default=".", help="directory scanned for *.json dictionaries")
    io.add_argument("--freq", default="frequency.txt", help="frequency list, most common first")
    io.add_argument("--out", default="steno_heatmap.png", help="output PNG path")
    io.add_argument("--glob", default="*.json", help="glob pattern for dictionary files")

    d = p.add_argument_group("data handling")
    d.add_argument("--length-mode", choices=["chars", "words"], default="chars",
                   help="measure word length by characters or by number of words")
    d.add_argument("--length-scale", choices=["log", "linear"], default="log",
                   help="warp the length (X) axis; log gives each doubling equal width")
    d.add_argument("--rarity-scale", choices=["log", "linear"], default="log",
                   help="warp the rarity (Y) axis; log gives each decade equal height")
    d.add_argument("--drop-unranked", action="store_true",
                   help="exclude words absent from the frequency list "
                        "(default: keep them, pinned to the rarest end)")
    d.add_argument("--keep-meta", action="store_true",
                   help="keep Plover formatting/meta entries ({...}, numbers, etc.)")
    d.add_argument("--aggregate", choices=["mean", "min", "median"], default="mean",
                   help="how to combine stroke counts of words sharing a cell")

    v = p.add_argument_group("appearance")
    v.add_argument("--max-strokes", type=int, default=5,
                   help="stroke count that maps to full red; higher counts clamp to it")
    v.add_argument("--grid", type=int, default=220, help="heat-map grid resolution")
    v.add_argument("--smooth", type=float, default=6.0,
                   help="Gaussian smoothing sigma in grid cells (0 = off / blocky)")
    v.add_argument("--shade", choices=["solid", "density", "mask"], default="solid",
                   help="solid: fill whole plot with full-opacity color (no black); "
                        "density: brightness follows word density; "
                        "mask: color only where words exist, black elsewhere")
    v.add_argument("--gamma", type=float, default=0.7,
                   help="brightness gamma for --shade density (lower = thin regions show more)")
    v.add_argument("--length-min", type=float, default=None,
                   help="left edge of length axis (default: shortest word)")
    v.add_argument("--length-max", type=float, default=None,
                   help="right edge of length axis (default: 99th percentile length)")
    v.add_argument("--scatter", action="store_true",
                   help="draw raw points instead of a smoothed heat-map")
    v.add_argument("--point-size", type=float, default=4.0, help="scatter point size")
    v.add_argument("--point-alpha", type=float, default=0.25, help="scatter point alpha")
    v.add_argument("--title", default="Strokes by word length * rarity",
                   help="chart title")
    v.add_argument("--width", type=float, default=11.0, help="figure width (inches)")
    v.add_argument("--height", type=float, default=8.5, help="figure height (inches)")
    v.add_argument("--dpi", type=int, default=140, help="output resolution")
    return p.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)

    out_abs = os.path.normcase(os.path.abspath(args.out))
    json_paths = [p for p in sorted(glob.glob(os.path.join(args.dir, args.glob)))
                  if os.path.normcase(os.path.abspath(p)) != out_abs]
    if not json_paths:
        sys.exit(f"ERROR: no files matching {args.glob!r} in {args.dir!r}")

    exact_rank, lower_rank, n_ranked = load_frequency(args.freq)

    include_unranked = not args.drop_unranked
    rarity_axis = {
        "scale": args.rarity_scale,
        "n_ranked": n_ranked,
        "max_rank": n_ranked + 1 if include_unranked else n_ranked,
        "include_unranked": include_unranked,
    }

    fewest_strokes = collect_strokes(json_paths, filter_meta=not args.keep_meta)
    x, y, strokes = build_points(fewest_strokes, exact_rank, lower_rank, rarity_axis,
                                 length_mode=args.length_mode)
    if x.size == 0:
        sys.exit("ERROR: nothing to plot. Check that --freq matches the dictionary "
                 "words, or drop --drop-unranked.")

    if args.length_min is None:
        args.length_min = float(max(1, int(np.min(x))))
    if args.length_max is None:
        args.length_max = float(int(np.ceil(np.percentile(x, 99))))
    if args.length_max <= args.length_min:
        args.length_max = args.length_min + 1

    length_axis = {
        "scale": args.length_scale,
        "lo": args.length_min,
        "hi": args.length_max,
        "mode": args.length_mode,
    }

    render = render_scatter if args.scatter else render_heatmap
    render(x, y, strokes, length_axis, rarity_axis, args)
    print(f"Saved -> {args.out}")


if __name__ == "__main__":
    main()