#!/usr/bin/env python3
"""RECLASSIFIER agent: re-classify all 34 sessions by reading actual content."""

import json
from collections import Counter
from pathlib import Path
from statistics import mean

OUTPUT_DIR = Path("/home/sagar/trace/docs/run1-09022026/analysis-outputs")
PARSED_DIR = OUTPUT_DIR / "parsed"
METRICS_FILE = OUTPUT_DIR / "metrics.json"

# Manual classifications based on reading all session content
CLASSIFICATIONS = {
    # 1. pip uninstall wandb, configuring driver.py setup
    "rollout-2026-01-28T15-51-17": {
        "new_type": "setup", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer uninstalling wandb and configuring driver.py project setup"
    },
    # 2. wandb API key configuration, login, env vars
    "rollout-2026-01-28T15-57-48": {
        "new_type": "setup", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Configuring wandb API key, login, and environment variables for experiment tracking"
    },
    # 3. 20 events, 9.6s, pasted driver output then interrupted immediately
    "rollout-2026-01-28T17-07-55": {
        "new_type": "context-only", "new_outcome": "context-only", "trivial": True,
        "confidence": "high",
        "reasoning": "Brief interrupted session - pasted driver.py output then immediately aborted"
    },
    # 4. Ground truth: context-only. 9 events, 0.1s
    "rollout-2026-01-28T17-08-14": {
        "new_type": "context-only", "new_outcome": "context-only", "trivial": True,
        "confidence": "high",
        "reasoning": "Only 9 events, 0.1s duration, no meaningful interaction occurred"
    },
    # 5. Ran driver.py, got apply_patch warning, interrupted
    "rollout-2026-01-28T17-08-38": {
        "new_type": "bug-fix", "new_outcome": "abandoned", "trivial": False,
        "confidence": "medium",
        "reasoning": "Developer ran driver.py, hit apply_patch tool error, session interrupted mid-fix"
    },
    # 6. "how can i improve my model?" + "Feature crosses: add explicit interactions"
    "rollout-2026-01-28T20-00-48": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer asked how to improve model then requested feature crosses and learned cross interactions"
    },
    # 7. "i changed NEG_PER_POS from 7 to 20, the test_loss changed from 0.12 to 0.08"
    "rollout-2026-01-28T21-21-52": {
        "new_type": "exploration", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer reporting experiment results (NEG_PER_POS change), exploring impact on test_loss"
    },
    # 8. 9 events, 0.2s, just config shown then nothing
    "rollout-2026-01-29T10-18-05": {
        "new_type": "context-only", "new_outcome": "context-only", "trivial": True,
        "confidence": "high",
        "reasoning": "Only 9 events, 0.2s duration, config file shown but no interaction"
    },
    # 9. Extensive code changes across config, model, training, ctr_ranker, driver
    "rollout-2026-01-29T10-18-14": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Major coding session touching config, model, training, ctr_ranker, and driver files with many apply_patch operations"
    },
    # 10. Ran driver.py showing wandb training run, then "."
    "rollout-2026-01-29T11-47-50": {
        "new_type": "exploration", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer ran training, pasted wandb output to check progress, minimal interaction ('.')"
    },
    # 11. 9 events, 0.4s, user said 'hi' then interrupted
    "rollout-2026-01-29T12-06-02": {
        "new_type": "context-only", "new_outcome": "context-only", "trivial": True,
        "confidence": "high",
        "reasoning": "User said 'hi' then immediately interrupted, 9 events in 0.4s"
    },
    # 12. 25 events, 6.8s, request truncated then interrupted
    "rollout-2026-01-29T13-21-42": {
        "new_type": "context-only", "new_outcome": "context-only", "trivial": True,
        "confidence": "high",
        "reasoning": "Very brief session, request didn't finish typing before interruption, 25 events in 6.8s"
    },
    # 13. Working on config, recommend, driver - "all the env f[lags]", "wire them in"
    "rollout-2026-01-29T13-23-37": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer adding environment flag support and wiring configuration into recommend and driver modules"
    },
    # 14. Data prep code - "is my [code correct]" then "do 1."
    "rollout-2026-01-29T20-07-43": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer working on data prep pipeline, verifying code then requesting implementation of option 1"
    },
    # 15. "can you fix error in model.py" + "can you fix error in training.py"
    "rollout-2026-01-29T20-30-24": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer explicitly requesting error fixes in model.py and training.py"
    },
    # 16. "can you fix error in training.py"
    "rollout-2026-01-29T20-37-57": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer requesting fix for error in training.py, single fix request"
    },
    # 17. Repeated test_loader selection in training.py
    "rollout-2026-01-29T20-44-10": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": False,
        "confidence": "medium",
        "reasoning": "Developer repeatedly selecting test_loader in training.py, iterating on fix for test data loading issue"
    },
    # 18. "fix my recommend.py"
    "rollout-2026-01-29T20-53-38": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer requesting fix for recommend.py, single request"
    },
    # 19. Pasted terminal output from running driver.py with progress bars + warning path
    "rollout-2026-01-29T22-46-03": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer pasted driver.py execution output showing a Python path warning, debugging runtime issue"
    },
    # 20. Ground truth: bug-fix. "this is my main fil[e]" + working on dataset/data_prep
    "rollout-2026-01-29T22-48-55": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer sharing main file context and working on dataset/data_prep bug fix"
    },
    # 21. Pasted terminal output from driver.py showing data prep progress
    "rollout-2026-01-29T23-07-18": {
        "new_type": "bug-fix", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer pasted driver.py output for debugging, terminal shows data preparation pipeline progress"
    },
    # 22. MODE='test' config selection, short session
    "rollout-2026-01-29T23-17-23": {
        "new_type": "setup", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer configuring MODE='test' setting, brief configuration change"
    },
    # 23. "exhausting permanent memory during write of data prep files" - asking about compressed formats
    "rollout-2026-01-30T13-43-01": {
        "new_type": "exploration", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer exploring solutions for memory exhaustion during data prep writes, asking about npz vs npy compressed formats"
    },
    # 24. 68 events, 177s, truncated request "(ni" then interrupted
    "rollout-2026-01-30T16-43-54": {
        "new_type": "context-only", "new_outcome": "abandoned", "trivial": True,
        "confidence": "medium",
        "reasoning": "Request truncated to '(ni', session interrupted before meaningful interaction, treated as abandoned context"
    },
    # 25. Continuation of training/dataset work, "(ni" then "yes"
    "rollout-2026-01-30T16-47-16": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer approving code changes to training.py and dataset.py, continuation of pipeline work"
    },
    # 26. "can recommend.py be made faster? how?" - then abandoned
    "rollout-2026-01-31T06-14-13": {
        "new_type": "refactor", "new_outcome": "abandoned", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer asked about recommend.py performance optimization but session was abandoned"
    },
    # 27. "can recommend.py be made faster? how?" - short exploration
    "rollout-2026-01-31T06-23-05": {
        "new_type": "exploration", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer asking how to optimize recommend.py, received analysis, short Q&A session"
    },
    # 28. Asked about speed + gave specific optimization instructions (precompute, topk, inference_mode)
    "rollout-2026-01-31T06-25-34": {
        "new_type": "refactor", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer gave specific optimization instructions: precompute tensors, replace batch topk, use torch.inference_mode()"
    },
    # 29. Iterative optimization of recommend.py - biggest win, GPU memory, auto-scaling
    "rollout-2026-01-31T06-30-50": {
        "new_type": "refactor", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Extended iterative optimization of recommend.py: GPU memory management, batch sizing, auto-scaling removal"
    },
    # 30. "how to make recommend.py faster?" + nvidia-smi + yes
    "rollout-2026-01-31T06-49-43": {
        "new_type": "refactor", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Developer optimizing recommend.py performance, checking GPU utilization with nvidia-smi"
    },
    # 31. "currently takes 8:47:41, want 3 hrs" - major perf optimization
    "rollout-2026-01-31T06-55-05": {
        "new_type": "refactor", "new_outcome": "resolved", "trivial": False,
        "confidence": "high",
        "reasoning": "Major performance optimization: recommend.py runtime from 8:47:41 to target 3hrs, touching data prep, config, mapping"
    },
    # 32. "this is from another repo" + pasting code from production recommendation system
    "rollout-2026-01-31T07-51-36": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": True,
        "confidence": "medium",
        "reasoning": "Developer importing code from another production repo to build new functionality"
    },
    # 33. "remove series already read by user, filter and recommend top 30"
    "rollout-2026-01-31T07-59-10": {
        "new_type": "feature", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer requesting new feature: filter out already-read series from recommendations, keep top 30"
    },
    # 34. "break recommend.py into proper functions" + remove backup values + config at top
    "rollout-2026-01-31T08-17-13": {
        "new_type": "refactor", "new_outcome": "resolved", "trivial": True,
        "confidence": "high",
        "reasoning": "Developer requesting code restructure: break into functions, remove fallback values, extract config to top of file"
    },
}


def main():
    # Load original metrics
    orig_metrics = json.loads(METRICS_FILE.read_text())

    # Build session lookup
    sessions_data = {}
    for f in sorted(PARSED_DIR.glob("*.json")):
        d = json.loads(f.read_text())
        sessions_data[d["filename"]] = d

    # Build reclassification
    reclass_sessions = []
    for d in orig_metrics["sessions"]:
        fn = d["filename"]
        # Find matching classification key
        key = None
        for k in CLASSIFICATIONS:
            if fn.startswith(k):
                key = k
                break

        if not key:
            print(f"WARNING: No classification for {fn}")
            continue

        c = CLASSIFICATIONS[key]
        reclass_sessions.append({
            "filename": fn,
            "old_type": d["type"],
            "new_type": c["new_type"],
            "old_outcome": d["outcome"],
            "new_outcome": c["new_outcome"],
            "trivial": c["trivial"],
            "confidence": c["confidence"],
            "reasoning": c["reasoning"],
        })

    # Summary
    new_type_counts = Counter(s["new_type"] for s in reclass_sessions)
    new_outcome_counts = Counter(s["new_outcome"] for s in reclass_sessions)
    trivial_count = sum(1 for s in reclass_sessions if s["trivial"])
    non_trivial_count = sum(1 for s in reclass_sessions if not s["trivial"])
    changes = sum(1 for s in reclass_sessions if s["old_type"] != s["new_type"] or s["old_outcome"] != s["new_outcome"])

    reclassification = {
        "sessions": reclass_sessions,
        "summary": {
            "by_type": {k: new_type_counts.get(k, 0) for k in ["bug-fix", "feature", "setup", "exploration", "refactor", "context-only"]},
            "by_outcome": {k: new_outcome_counts.get(k, 0) for k in ["resolved", "partial", "abandoned", "context-only"]},
            "trivial_count": trivial_count,
            "non_trivial_count": non_trivial_count,
            "changes_from_original": changes,
        }
    }

    reclass_file = OUTPUT_DIR / "reclassification.json"
    reclass_file.write_text(json.dumps(reclassification, indent=2))
    print(f"Wrote {reclass_file}")
    print(f"  Type distribution: {dict(new_type_counts)}")
    print(f"  Outcome distribution: {dict(new_outcome_counts)}")
    print(f"  Trivial: {trivial_count}, Non-trivial: {non_trivial_count}")
    print(f"  Changes from original: {changes}")

    # Write metrics-v2.json
    # Rebuild metrics with corrected classifications
    date_counts = Counter()
    tod_counts = Counter()
    durations = []
    turns_list = []
    specificities = []
    struggles = []
    sessions_v2 = []

    for s in reclass_sessions:
        fn = s["filename"]
        sd = sessions_data[fn]
        date_counts[sd["date"]] += 1
        # Recompute time_of_day from parsed data
        tod = sd.get("metrics", {}).get("time_of_day", "unknown")
        tod_counts[tod] += 1
        durations.append(sd["duration_seconds"])
        turns_list.append(sd["turn_count"])
        specificities.append(sd.get("metrics", {}).get("prompt_specificity", 1))

        # Recompute struggle score with new outcome
        score = 0
        tc = sd["turn_count"]
        if tc > 10:
            score += 2
        if tc > 20:
            score += 2
        user_texts = [m.get("text", "").lower().strip() for m in sd.get("user_messages", [])]
        if len(user_texts) >= 3:
            for i in range(len(user_texts)):
                for j in range(i + 1, len(user_texts)):
                    if user_texts[i] and user_texts[j] and (
                        user_texts[i] == user_texts[j] or
                        (len(user_texts[i]) > 20 and user_texts[i] in user_texts[j]) or
                        (len(user_texts[j]) > 20 and user_texts[j] in user_texts[i])
                    ):
                        score += 2
                        break
                else:
                    continue
                break
        error_msg_count = sum(1 for t in user_texts if any(w in t for w in ["error", "fail"]))
        if error_msg_count >= 2:
            score += 2
        if s["new_outcome"] == "abandoned":
            score += 2
        if sd["duration_seconds"] > 600 and s["new_outcome"] != "resolved":
            score += 2
        score = min(score, 10)
        struggles.append(score)

        sessions_v2.append({
            "filename": fn,
            "outcome": s["new_outcome"],
            "type": s["new_type"],
            "turns": sd["turn_count"],
            "duration": sd["duration_seconds"],
            "specificity": sd.get("metrics", {}).get("prompt_specificity", 1),
            "struggle": score,
            "trivial": s["trivial"],
        })

    metrics_v2 = {
        "total_sessions": len(reclass_sessions),
        "by_outcome": {k: new_outcome_counts.get(k, 0) for k in ["resolved", "abandoned", "partial", "context-only"]},
        "by_type": {k: new_type_counts.get(k, 0) for k in ["bug-fix", "feature", "setup", "exploration", "refactor", "context-only"]},
        "avg_turns": round(mean(turns_list), 1) if turns_list else 0,
        "avg_duration_seconds": round(mean(durations), 1) if durations else 0,
        "avg_specificity": round(mean(specificities), 1) if specificities else 0,
        "avg_struggle_score": round(mean(struggles), 1) if struggles else 0,
        "by_date": dict(sorted(date_counts.items())),
        "by_time_of_day": {k: tod_counts.get(k, 0) for k in ["morning", "afternoon", "evening", "night"]},
        "trivial_count": trivial_count,
        "non_trivial_count": non_trivial_count,
        "sessions": sessions_v2,
    }

    metrics_v2_file = OUTPUT_DIR / "metrics-v2.json"
    metrics_v2_file.write_text(json.dumps(metrics_v2, indent=2))
    print(f"Wrote {metrics_v2_file}")


if __name__ == "__main__":
    main()
