#!/usr/bin/env python3
"""
Track per-parameter direction across a sequence of .dna files (e.g.
one KEPT candidate.dna per cycle.sh cycle, in chronological order) to
tell apart two very different things that both show up as "the DNA
changed again":

  - a parameter trending: consecutive deltas keep the same sign, i.e.
    the tuner keeps pushing it the same direction cycle over cycle --
    this looks like real signal.
  - a parameter flip-flopping: consecutive deltas alternate sign --
    the tuner is chasing sampling noise in the position batch, not
    converging on anything.

For each raw DNA cell we compute:
    net    = final_value - first_value
    churn  = sum(|delta| for each consecutive step)
    consistency = net / churn   (in [-1, 1]; 0 churn -> consistency 1 if
                                  net is also 0, else undefined/skipped)

consistency near +-1 means every step moved the same direction (pure
trend); consistency near 0 with nonzero churn means it moved a lot but
ended up roughly where it started (pure flip-flop).

Usage:
    python3 dna_trend.py cycle1.dna cycle2.dna cycle3.dna ...
    python3 dna_trend.py --cycles-dir eval_tune/cycles   # auto-discover,
        chronological by directory timestamp, only cycles with a KEPT
        candidate.dna per their summary.txt

Prints, per named array, a one-line summary, then the individual cells
with the strongest trend and the worst flip-flop for a closer look.
"""
import argparse
import sys
from pathlib import Path

from dna_diff import DNA_NAMES, read_dna_file


def load_sequence(paths):
    rows_by_file = [read_dna_file(p) for p in paths]
    n_cells_per_file = [sum(len(row) for row in rows) for rows in rows_by_file]
    if len(set(n_cells_per_file)) != 1:
        sys.exit(f"cell-count mismatch across files: {dict(zip(paths, n_cells_per_file))} "
                  f"-- these .dna files don't all match the same eval.c revision.")
    # Flatten each file to one list of (name, index_within_array, value)
    flat_sequences = []
    for rows in rows_by_file:
        flat = []
        for name, row in zip(DNA_NAMES, rows):
            for i, v in enumerate(row):
                flat.append((name, i, v))
        flat_sequences.append(flat)
    return flat_sequences


def discover_kept_cycles(cycles_dir):
    paths = []
    for d in sorted(Path(cycles_dir).iterdir()):
        summary = d / "summary.txt"
        candidate = d / "candidate.dna"
        if summary.exists() and candidate.exists():
            text = summary.read_text()
            if "RESULT: KEPT" in text:
                paths.append(candidate)
    return paths


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dna_files", nargs="*")
    ap.add_argument("--cycles-dir", default=None,
                     help="auto-discover KEPT candidate.dna files under this "
                          "cycle.sh cycles/ directory, chronologically")
    ap.add_argument("--top", type=int, default=15,
                     help="how many strongest-trend / worst-flip-flop cells to list")
    args = ap.parse_args()

    if args.cycles_dir:
        paths = discover_kept_cycles(args.cycles_dir)
    else:
        paths = [Path(p) for p in args.dna_files]

    if len(paths) < 2:
        sys.exit("need at least 2 .dna files (in chronological order) to "
                  "compute a trend -- got "
                  f"{len(paths)}: {[str(p) for p in paths]}")

    print(f"Sequence ({len(paths)} points, chronological):")
    for p in paths:
        print(f"  {p}")
    print()

    sequences = load_sequence(paths)
    n_cells = len(sequences[0])

    results = []  # (name, index, net, churn, consistency, first, last)
    for cell_idx in range(n_cells):
        name, arr_idx, _ = sequences[0][cell_idx]
        values = [seq[cell_idx][2] for seq in sequences]
        deltas = [b - a for a, b in zip(values, values[1:])]
        net = values[-1] - values[0]
        churn = sum(abs(d) for d in deltas)
        if churn == 0:
            continue  # never moved -- not interesting either way
        consistency = net / churn
        results.append((name, arr_idx, net, churn, consistency, values[0], values[-1]))

    if not results:
        print("No cell changed at all across this sequence.")
        return

    # Per-array rollup: mean |consistency| weighted by churn, plus counts.
    by_array = {}
    for name, arr_idx, net, churn, consistency, first, last in results:
        d = by_array.setdefault(name, {"churn": 0, "weighted": 0.0, "n": 0,
                                        "trending": 0, "flipping": 0})
        d["churn"] += churn
        d["weighted"] += abs(consistency) * churn
        d["n"] += 1
        if abs(consistency) >= 0.6:
            d["trending"] += 1
        elif abs(consistency) <= 0.25:
            d["flipping"] += 1

    print(f"{'ARRAY':45s} {'cells':>6s} {'trend':>6s} {'flip':>5s} {'churn-wtd consistency':>22s}")
    for name in DNA_NAMES:
        if name not in by_array:
            continue
        d = by_array[name]
        wavg = d["weighted"] / d["churn"] if d["churn"] else 0.0
        print(f"{name:45s} {d['n']:6d} {d['trending']:6d} {d['flipping']:5d} {wavg:22.2f}")

    results.sort(key=lambda r: -abs(r[4]) * r[3])  # weight by churn too
    trending = [r for r in results if r[4] >= 0.6][:args.top]
    flipping = sorted([r for r in results if abs(r[4]) <= 0.25],
                       key=lambda r: -r[3])[:args.top]

    print(f"\nTop {len(trending)} most consistently-trending cells "
          f"(same direction every step):")
    for name, idx, net, churn, cons, first, last in trending:
        print(f"  {name}[{idx}]: {first} -> {last} "
              f"(net {net:+d}, churn {churn}, consistency {cons:+.2f})")

    print(f"\nTop {len(flipping)} worst flip-flopping cells "
          f"(moved a lot, net ~0 -- likely noise):")
    for name, idx, net, churn, cons, first, last in flipping:
        print(f"  {name}[{idx}]: {first} -> {last} "
              f"(net {net:+d}, churn {churn}, consistency {cons:+.2f})")


if __name__ == "__main__":
    main()