"""parse_v6_log.py — Parse v6_progressive.log and extract evolution data.

Extracts stage results from the log file since training was killed before
writing the final report.
"""
import json
import re
from pathlib import Path
from datetime import datetime

LOG_PATH = Path("/home/z/my-project/logs/v6_progressive.log")
REPORT_PATH = Path("/home/z/my-project/BiGRU_T_version/v6_progressive_report.json")

# Regex patterns
STAGE_RE = re.compile(r"STAGE: (\w+) \(n=(\d+), epoch=(\d+)\)")
STAGE_END_RE = re.compile(
    r"STAGE END \| val_loss=([\d.]+) \| train_avg=([\d.]+) \| ppl=([\d.]+) "
    r"\| gap=([+-]?[\d.]+) \| eval_acc=([\d.]+)% \| elapsed=([\d.]+)s \| n_steps=(\d+)"
)
SOM_RE = re.compile(r"SOM: σ=([\d.]+), α=([\d.]+), step=(\d+)")
STEP_RE = re.compile(
    r"step\s+(\d+)\s+\|\s+loss=([\d.]+)\s+\|\s+tps=(\d+)\s+\|\s+tflops=([\d.]+)\s+\|\s+\|\|g\|\|=([\d.]+)\s+\|\s+cpu=(\d+)%\s+\|\s+rss=(\d+)MB"
)


def parse_log():
    """Parse the log file and return evolution data."""
    if not LOG_PATH.exists():
        print(f"[ERROR] Log file not found: {LOG_PATH}")
        return None

    with open(LOG_PATH) as f:
        lines = f.readlines()

    print(f"[INFO] Parsing {len(lines)} log lines from {LOG_PATH}")

    evolution = []
    current_stage = None
    current_n = None
    current_epoch = None
    current_steps = []
    som_data = {}

    for line in lines:
        # STAGE start
        m = STAGE_RE.search(line)
        if m:
            # Save previous stage if any
            if current_stage is not None:
                # Compute summary
                if current_steps:
                    losses = [s["loss"] for s in current_steps]
                    tps_arr = [s["tps"] for s in current_steps]
                    tflops_arr = [s["tflops"] for s in current_steps]
                    rss_arr = [s["rss_mb"] for s in current_steps]
                    avg_tflops = sum(tflops_arr) / len(tflops_arr)
                    avg_tps = sum(tps_arr) / len(tps_arr)
                    final_rss = rss_arr[-1]
                    max_rss = max(rss_arr)
                else:
                    avg_tflops = 0
                    avg_tps = 0
                    final_rss = 0
                    max_rss = 0
                # Look for STAGE END
                evolution.append({
                    "stage": current_stage,
                    "n_samples": current_n,
                    "epoch": current_epoch,
                    "n_steps": len(current_steps),
                    "avg_tflops": avg_tflops,
                    "avg_tokens_per_sec": avg_tps,
                    "final_rss_mb": final_rss,
                    "max_rss_mb": max_rss,
                    "train_loss_mean": sum(losses) / len(losses) if current_steps else 0,
                    "val_loss": None,  # filled below
                    "perplexity": None,
                    "overfitting_gap": None,
                    "eval_acc": None,
                    "elapsed_s": None,
                    "som": dict(som_data),
                })
                som_data = {}
            current_stage = m.group(1)
            current_n = int(m.group(2))
            current_epoch = int(m.group(3))
            current_steps = []
            continue

        # STAGE END
        m = STAGE_END_RE.search(line)
        if m and current_stage is not None:
            val_loss = float(m.group(1))
            train_avg = float(m.group(2))
            ppl = float(m.group(3))
            gap = float(m.group(4))
            eval_acc = float(m.group(5)) / 100.0
            elapsed = float(m.group(6))
            n_steps_end = int(m.group(7))
            # Update last entry's stage-end info
            # Actually, save as a new "complete" entry
            # If evolution has an entry for this stage without val_loss, update it
            if evolution and evolution[-1]["stage"] == current_stage and evolution[-1]["n_samples"] == current_n and evolution[-1]["epoch"] == current_epoch:
                evolution[-1].update({
                    "val_loss": val_loss,
                    "train_loss_mean": train_avg,
                    "perplexity": ppl,
                    "overfitting_gap": gap,
                    "eval_acc": eval_acc,
                    "elapsed_s": elapsed,
                    "n_steps": n_steps_end,
                })
            else:
                evolution.append({
                    "stage": current_stage,
                    "n_samples": current_n,
                    "epoch": current_epoch,
                    "n_steps": n_steps_end,
                    "avg_tflops": 0,
                    "avg_tokens_per_sec": 0,
                    "final_rss_mb": 0,
                    "max_rss_mb": 0,
                    "train_loss_mean": train_avg,
                    "val_loss": val_loss,
                    "perplexity": ppl,
                    "overfitting_gap": gap,
                    "eval_acc": eval_acc,
                    "elapsed_s": elapsed,
                    "som": dict(som_data),
                })
                som_data = {}
            continue

        # SOM
        m = SOM_RE.search(line)
        if m and current_stage is not None:
            som_data = {
                "sigma": float(m.group(1)),
                "alpha": float(m.group(2)),
                "step": int(m.group(3)),
            }
            continue

        # STEP
        m = STEP_RE.search(line)
        if m and current_stage is not None:
            current_steps.append({
                "step": int(m.group(1)),
                "loss": float(m.group(2)),
                "tps": int(m.group(3)),
                "tflops": float(m.group(4)),
                "grad_norm": float(m.group(5)),
                "cpu_pct": int(m.group(6)),
                "rss_mb": int(m.group(7)),
            })
            continue

    # Save last stage if incomplete (no STAGE END)
    if current_stage is not None and current_steps:
        losses = [s["loss"] for s in current_steps]
        tps_arr = [s["tps"] for s in current_steps]
        tflops_arr = [s["tflops"] for s in current_steps]
        rss_arr = [s["rss_mb"] for s in current_steps]
        evolution.append({
            "stage": current_stage,
            "n_samples": current_n,
            "epoch": current_epoch,
            "n_steps": len(current_steps),
            "avg_tflops": sum(tflops_arr) / len(tflops_arr) if tflops_arr else 0,
            "avg_tokens_per_sec": sum(tps_arr) / len(tps_arr) if tps_arr else 0,
            "final_rss_mb": rss_arr[-1] if rss_arr else 0,
            "max_rss_mb": max(rss_arr) if rss_arr else 0,
            "train_loss_mean": sum(losses) / len(losses) if losses else 0,
            "val_loss": None,
            "perplexity": None,
            "overfitting_gap": None,
            "eval_acc": None,
            "elapsed_s": None,
            "som": dict(som_data),
            "incomplete": True,
        })

    # Filter out empty stages (no steps, no end)
    evolution = [e for e in evolution if e.get("n_steps", 0) > 0 or e.get("val_loss") is not None]

    # Merge duplicate entries (same stage+n+epoch): one has step stats, other has val_loss
    merged = {}
    order = []
    for e in evolution:
        key = (e["stage"], e["n_samples"], e["epoch"])
        if key not in merged:
            merged[key] = {}
            order.append(key)
        # Merge — later entries overwrite None
        for k, v in e.items():
            if v is not None and v != 0:
                merged[key][k] = v
            elif k not in merged[key]:
                merged[key][k] = v
        # Ensure all keys present
        for k in ["stage", "n_samples", "epoch", "n_steps", "avg_tflops", "avg_tokens_per_sec",
                  "final_rss_mb", "max_rss_mb", "train_loss_mean", "val_loss", "perplexity",
                  "overfitting_gap", "eval_acc", "elapsed_s", "som"]:
            if k not in merged[key]:
                merged[key][k] = None
        # Mark incomplete (no val_loss)
        if merged[key].get("val_loss") is None:
            merged[key]["incomplete"] = True

    evolution = [merged[k] for k in order]

    print(f"[INFO] Parsed {len(evolution)} stage entries (after merge)")
    return evolution


def main():
    evolution = parse_log()
    if not evolution:
        print("[ERROR] No evolution data parsed")
        return

    # Print summary
    print(f"\n{'=' * 100}")
    print(f"  EVOLUÇÃO (parsed from log)")
    print(f"{'=' * 100}")
    print(f"  {'#':<3} {'stage':<14} {'n':>4} {'ep':>3} {'val_loss':>10} {'ppl':>10} {'eval_acc':>10} {'rss_mb':>8} {'tflops':>8} {'incomp':>7}")
    print(f"  {'-'*3} {'-'*14} {'-'*4} {'-'*3} {'-'*10} {'-'*10} {'-'*10} {'-'*8} {'-'*8} {'-'*7}")
    for i, e in enumerate(evolution):
        val = f"{e['val_loss']:.4f}" if e.get("val_loss") is not None else "—"
        ppl = f"{e['perplexity']:.2f}" if e.get("perplexity") is not None else "—"
        acc = f"{e['eval_acc']*100:.2f}%" if e.get("eval_acc") is not None else "—"
        rss = f"{e['final_rss_mb']:.0f}" if e.get("final_rss_mb") else "—"
        tf = f"{e['avg_tflops']:.4f}" if e.get("avg_tflops") else "—"
        incomp = "YES" if e.get("incomplete") else ""
        print(f"  {i+1:<3} {e['stage']:<14} {e['n_samples']:>4} {e['epoch']:>3} "
              f"{val:>10} {ppl:>10} {acc:>10} {rss:>8} {tf:>8} {incomp:>7}")

    # Save partial report
    report = {
        "version": "V6_progressive_partial",
        "timestamp": datetime.now().isoformat(),
        "note": "Partial report — training was killed before full completion. Evolution data extracted from log.",
        "config": {
            "epochs": 1,
            "batch_size": 8,
            "max_seq_len": 24,
            "num_layers_hyp": 8,
            "stages_punctuation": [100, 200, 300, 400, 500],
            "stages_text": [100, 200, 300, 400, 500],
            "stages_heavy": [10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
        },
        "evolution": evolution,
        "n_stages_captured": len(evolution),
        "n_stages_total_expected": 20,  # 5 + 5 + 10
        "completion_pct": len(evolution) / 20 * 100,
    }
    REPORT_PATH.write_text(json.dumps(report, indent=2, ensure_ascii=False, default=str))
    print(f"\n[INFO] Partial report saved: {REPORT_PATH}")
    print(f"[INFO] Captured {len(evolution)}/20 stages ({len(evolution)/20*100:.1f}%)")


if __name__ == "__main__":
    main()