import json
from collections import defaultdict
from pathlib import Path
from statistics import mean, median, variance
from typing import Any, Dict, List, Optional


PROCESSED_DIR_NAME = "processed"
METRICS_PER_PLAYER_DIR_NAME = "metrics_per_player"
GROUPED_OUTPUT_FILE = "metrics_summary_grouped.json"
OVERALL_OUTPUT_FILE = "metrics_summary_overall.json"
SUMMARY_DIR_NAME = "metrics_summary"


Numeric = float


def _load_player_metrics(metrics_dir: Path) -> List[Dict[str, Any]]:
    if not metrics_dir.exists() or not metrics_dir.is_dir():
        raise FileNotFoundError(f"Metrics directory not found: {metrics_dir}")

    players: List[Dict[str, Any]] = []
    for path in sorted(metrics_dir.glob("*.json")):
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        players.append(data)
    if not players:
        raise ValueError(f"No metrics JSON files found in {metrics_dir}")
    return players


def _get_quiz_value(quiz_answers: List[Dict[str, Any]], identifier: str) -> Optional[str]:
    for ans in quiz_answers:
        if ans.get("identifier") == identifier:
            value = ans.get("value")
            if value is None:
                return None
            return str(value)
    return None


def _normalize_continue(value: Optional[str]) -> Optional[bool]:
    if value is None:
        return None
    v = value.strip().lower()
    if v in {"yes", "y", "true", "1"}:
        return True
    if v in {"no", "n", "false", "0"}:
        return False
    return None


def _normalize_challenge(value: Optional[str]) -> Optional[str]:
    if value is None:
        return None
    v = value.strip().lower().replace(" ", "")

    # Map various possible wordings into canonical keys
    if "too" in v and "easy" in v:
        return "tooeasy"
    if "too" in v and ("difficult" in v or "hard" in v):
        return "toodifficult"
    if "optimal" in v or "justright" in v or "justrightdifficulty" in v:
        return "optimal"

    return None


def _collect_numeric(player: Dict[str, Any], key: str) -> Optional[Numeric]:
    value = player.get(key)
    if isinstance(value, (int, float)):
        return float(value)
    try:
        if value is not None:
            return float(value)
    except (TypeError, ValueError):
        return None
    return None


def _compute_stats(values: List[Numeric]) -> Optional[Dict[str, Numeric]]:
    if not values:
        return None

    stats: Dict[str, Numeric] = {
        "mean": float(mean(values)),
        "median": float(median(values)),
    }

    # Sample variance is only defined for n >= 2
    if len(values) >= 2:
        stats["variance"] = float(variance(values))
    else:
        stats["variance"] = None  # type: ignore[assignment]

    return stats


def _summarize_lz_window_metrics(players: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Aggregate LZ metrics per level, treating all windows of a single level
    for a given player as a whole.

    Implementation detail:
    - For each player and level, first aggregate that player's windows by
      computing the mean lz_phrase_count, lz_normalized_complexity and apm
      over their windows for that level.
    - Then, for each level, compute summary stats over the per-player
      aggregated values (so each player contributes at most once per level).
    """
    # level -> list of (player_level_mean_phrase, player_level_mean_complexity,
    #                   player_level_mean_apm, player_level_window_count)
    level_player_metrics: Dict[str, List[tuple[Numeric, Numeric, Numeric, int]]] = {}

    for player in players:
        windows = player.get("lz_window_metrics")
        if not isinstance(windows, list):
            continue

        # Group this player's windows by level
        player_level_windows: Dict[str, List[tuple[Numeric, Numeric, Numeric]]] = {}
        for window in windows:
            if not isinstance(window, dict):
                continue

            level = window.get("level")
            if level is None:
                continue

            level_key = str(level)
            lz_phrase_count = window.get("lz_phrase_count")
            lz_complexity = window.get("lz_normalized_complexity")
            apm = window.get("apm")

            if (
                isinstance(lz_phrase_count, (int, float))
                and isinstance(lz_complexity, (int, float))
                and isinstance(apm, (int, float))
            ):
                if level_key not in player_level_windows:
                    player_level_windows[level_key] = []
                player_level_windows[level_key].append(
                    (
                        float(lz_phrase_count),
                        float(lz_complexity),
                        float(apm),
                    )
                )

        # Collapse this player's windows into a single sample per level
        for level_key, metrics_list in player_level_windows.items():
            if not metrics_list:
                continue

            phrase_values = [m[0] for m in metrics_list]
            complexity_values = [m[1] for m in metrics_list]
            apm_values = [m[2] for m in metrics_list]

            mean_phrase = float(mean(phrase_values))
            mean_complexity = float(mean(complexity_values))
            mean_apm = float(mean(apm_values))
            window_count = len(metrics_list)

            if level_key not in level_player_metrics:
                level_player_metrics[level_key] = []
            level_player_metrics[level_key].append(
                (mean_phrase, mean_complexity, mean_apm, window_count)
            )

    per_level: Dict[str, Any] = {}
    for level_key, metrics_list in level_player_metrics.items():
        if not metrics_list:
            continue

        # One aggregated value per player-level
        phrase_means = [m[0] for m in metrics_list]
        complexity_means = [m[1] for m in metrics_list]
        apm_means = [m[2] for m in metrics_list]

        phrase_stats = _compute_stats(phrase_means)
        complexity_stats = _compute_stats(complexity_means)
        apm_stats = _compute_stats(apm_means)

        total_windows = sum(int(m[3]) for m in metrics_list)

        per_level[level_key] = {
            # Total number of underlying windows that contributed to this level
            "window_count": total_windows,
            "lz_phrase_count": phrase_stats if phrase_stats else {
                "mean": None,
                "median": None,
                "variance": None,
            },
            "apm": apm_stats if apm_stats else {
                "mean": None,
                "median": None,
                "variance": None,
            },
            "lz_normalized_complexity": complexity_stats if complexity_stats else {
                "mean": None,
                "median": None,
                "variance": None,
            },
        }

    return {"per_level": per_level}


def _summarize_group(players: List[Dict[str, Any]]) -> Dict[str, Any]:
    metric_keys = [
        "level_reached",
        "average_apm",
        "lz_phrase_count",
        "lz_average_phrase_length",
        "lz_normalized_complexity",
    ]

    summary: Dict[str, Any] = {
        "player_count": len(players),
        "metrics": {},
    }

    for key in metric_keys:
        values: List[Numeric] = []
        for player in players:
            num = _collect_numeric(player, key)
            if num is not None:
                values.append(num)

        stats = _compute_stats(values)
        if stats is not None:
            summary["metrics"][key] = stats
        else:
            summary["metrics"][key] = {
                "mean": None,
                "median": None,
                "variance": None,
            }

    # Add aggregated LZ window metrics per level
    summary["lz_window_metrics"] = _summarize_lz_window_metrics(players)

    return summary


def compute_grouped_and_overall_summaries() -> Dict[str, Path]:
    script_dir = Path(__file__).parent
    processed_dir = script_dir / PROCESSED_DIR_NAME
    metrics_dir = processed_dir / METRICS_PER_PLAYER_DIR_NAME
    summary_dir = processed_dir / SUMMARY_DIR_NAME
    summary_dir.mkdir(parents=True, exist_ok=True)

    players = _load_player_metrics(metrics_dir)

    grouped_players: Dict[str, Dict[str, List[Dict[str, Any]]]] = {
        "would_continue": {
            "tooeasy": [],
            "toodifficult": [],
            "optimal": [],
        },
        "would_not_continue": {
            "tooeasy": [],
            "toodifficult": [],
            "optimal": [],
        },
    }

    for player in players:
        quiz_answers = player.get("quiz_answers") or []
        if not isinstance(quiz_answers, list):
            continue

        continue_raw = _get_quiz_value(quiz_answers, "continue")
        challenge_raw = _get_quiz_value(quiz_answers, "challenge")

        would_continue = _normalize_continue(continue_raw)
        challenge = _normalize_challenge(challenge_raw)

        if would_continue is None or challenge is None:
            # Skip players with missing or unrecognized quiz data for grouped stats,
            # but they will still contribute to overall stats.
            continue

        top_key = "would_continue" if would_continue else "would_not_continue"
        if challenge in grouped_players[top_key]:
            grouped_players[top_key][challenge].append(player)

    grouped_summary: Dict[str, Dict[str, Any]] = {
        "would_continue": {},
        "would_not_continue": {},
    }

    for top_key in grouped_players:
        for challenge_key, plist in grouped_players[top_key].items():
            grouped_summary[top_key][challenge_key] = _summarize_group(plist)

    overall_summary = _summarize_group(players)

    grouped_output_path = processed_dir / GROUPED_OUTPUT_FILE
    overall_output_path = processed_dir / OVERALL_OUTPUT_FILE

    # Original combined files in /processed
    with grouped_output_path.open("w", encoding="utf-8") as f:
        json.dump(grouped_summary, f, indent=2)

    with overall_output_path.open("w", encoding="utf-8") as f:
        json.dump(overall_summary, f, indent=2)

    # Additional per-group files in /processed and subfolders
    # Overall (no grouping)
    overall_summary_path = summary_dir / "overall.json"
    with overall_summary_path.open("w", encoding="utf-8") as f:
        json.dump(overall_summary, f, indent=2)

    # Per combination: continue/not-continue x difficulty
    per_group_paths: Dict[str, Path] = {}
    for top_key, challenge_dict in grouped_summary.items():
        group_dir = summary_dir / top_key
        group_dir.mkdir(parents=True, exist_ok=True)
        for challenge_key, summary in challenge_dict.items():
            file_path = group_dir / f"{challenge_key}.json"
            with file_path.open("w", encoding="utf-8") as f:
                json.dump(summary, f, indent=2)
            per_group_paths[f"{top_key}/{challenge_key}"] = file_path

    print(f"Grouped metrics summary written to: {grouped_output_path}")
    print(f"Overall metrics summary written to: {overall_output_path}")
    print(f"Detailed summaries written under: {summary_dir}")

    result_paths: Dict[str, Path] = {
        "grouped": grouped_output_path,
        "overall": overall_output_path,
        "overall_detailed": overall_summary_path,
    }
    result_paths.update(per_group_paths)

    return result_paths


if __name__ == "__main__":
    compute_grouped_and_overall_summaries()


