"""Main experiment orchestrator — runs the full A-to-Z evaluation pipeline."""

from __future__ import annotations

import asyncio
import json
import logging
import random
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

import numpy as np

from .cache import ResultCache
from .config import EXPERIMENT_CONFIG, RESULTS_DIR
from .corpus.loader import TestText, load_corpus
from .conditions.naked import NakedRunner
from .conditions.vocabulary_only import VocabularyOnlyRunner
from .conditions.llm_scaffold import LLMScaffoldRunner
from .judges.llm_client import JudgeLLM
from .judges.layer1_programmatic import evaluate_layer1
from .judges.layer2_lens_depth import evaluate_layer2
from .judges.layer3_integration import evaluate_layer3
from .judges.layer4_discovery import evaluate_layer4
from .judges.pairwise import run_all_pairwise
from .sanitize import sanitize_for_judge
from .stats.bradley_terry import fit_bradley_terry, strengths_to_elo
from .stats.hypothesis_tests import (
    cliffs_delta,
    test_all_hypotheses,
    wilcoxon_paired,
)
from .report import generate_report

logger = logging.getLogger(__name__)

CONDITIONS: List[str] = EXPERIMENT_CONFIG["conditions"]


class Experiment:
    """Full evaluation experiment."""

    def __init__(self, results_dir: Path | None = None):
        self.results_dir = results_dir or RESULTS_DIR
        self.results_dir.mkdir(parents=True, exist_ok=True)
        self.cache = ResultCache(self.results_dir / "cache")
        self.judge = JudgeLLM()
        self.runners: Dict[str, Any] = {
            "NAKED": NakedRunner(self.judge),
            "VOCAB_ONLY": VocabularyOnlyRunner(self.judge),
            "LLM_SCAFFOLD": LLMScaffoldRunner(self.judge),
        }

    # ── public API ────────────────────────────────────────────────

    async def run(
        self,
        skip_generation: bool = False,
        skip_judging: bool = False,
        max_texts: int | None = None,
    ) -> Path:
        """Execute the 7-phase experiment pipeline."""
        random.seed(EXPERIMENT_CONFIG["seed"])
        np.random.seed(EXPERIMENT_CONFIG["seed"])

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir = self.results_dir / f"run_{timestamp}"
        run_dir.mkdir(parents=True, exist_ok=True)

        # Phase 1 — corpus
        logger.info("Phase 1: Loading corpus …")
        corpus = load_corpus()
        if max_texts and max_texts < len(corpus):
            # Balanced subset: sample proportionally from each domain
            from collections import defaultdict
            by_domain = defaultdict(list)
            for t in corpus:
                by_domain[t.domain].append(t)
            per_domain = max(1, max_texts // len(by_domain))
            subset = []
            for domain in sorted(by_domain):
                subset.extend(by_domain[domain][:per_domain])
            corpus = subset[:max_texts]
        logger.info("  Loaded %d texts across %d domains",
                     len(corpus), len({t.domain for t in corpus}))

        # Phase 2 — generation
        if not skip_generation:
            logger.info("Phase 2: Generating outputs (%d runs) …",
                         len(corpus) * len(CONDITIONS))
            outputs = await self._generate_all(corpus, run_dir)
        else:
            logger.info("Phase 2: Loading cached outputs …")
            outputs = self._load_cached_outputs(run_dir)

        # Phase 3 — 4-layer evaluation
        if not skip_judging:
            logger.info("Phase 3: Evaluating outputs (4 layers) …")
            scores = await self._evaluate_all(corpus, outputs, run_dir)
        else:
            logger.info("Phase 3: Loading cached scores …")
            scores = self._load_cached_scores(run_dir)

        # Phase 4 — pairwise
        if not skip_judging:
            logger.info("Phase 4: Pairwise comparisons …")
            pairwise = await self._pairwise_all(corpus, outputs, run_dir)
        else:
            pairwise = self._load_cached_pairwise(run_dir)

        # Phase 5 — Bradley-Terry
        logger.info("Phase 5: Fitting Bradley-Terry model …")
        bt_records = self._pairwise_to_records(pairwise)
        bt_strengths = fit_bradley_terry(bt_records, CONDITIONS)
        elo_ratings = strengths_to_elo(bt_strengths)

        # Phase 6 — hypothesis tests
        logger.info("Phase 6: Running hypothesis tests …")
        condition_scores = self._aggregate_scores(scores)
        hypotheses = test_all_hypotheses(condition_scores)

        # H7/H8: difficulty and domain moderation (for LLM_SCAFFOLD)
        if "LLM_SCAFFOLD" in self.runners:
            hypotheses["H7_llm_scaffold"] = self._test_h7(
                corpus, scores, cond="LLM_SCAFFOLD",
            )
            hypotheses["H8_llm_scaffold"] = self._test_h8(
                corpus, scores, cond="LLM_SCAFFOLD",
            )

        # Phase 7 — report
        logger.info("Phase 7: Generating report …")
        report_path = generate_report(
            run_dir=run_dir,
            corpus=corpus,
            scores=scores,
            pairwise=pairwise,
            bt_strengths=bt_strengths,
            elo_ratings=elo_ratings,
            hypotheses=hypotheses,
        )
        logger.info("DONE — report at %s", report_path)
        return report_path

    # ── Phase 2 helpers ───────────────────────────────────────────

    async def _generate_all(
        self, corpus: List[TestText], run_dir: Path,
    ) -> Dict[str, Dict[str, Any]]:
        outputs: Dict[str, Dict[str, Any]] = {}
        for text in corpus:
            outputs[text.id] = {}
            for cond_id, runner in self.runners.items():
                cache_key = f"{text.id}_{cond_id}"
                cached = self.cache.get(cache_key)
                if cached:
                    from .conditions.base import ConditionOutput
                    outputs[text.id][cond_id] = ConditionOutput(**cached)
                    continue

                logger.info("  Generating: %s × %s", text.id, cond_id)
                result = await runner.run(text.text, text.id)
                outputs[text.id][cond_id] = result
                self.cache.set(cache_key, {
                    "condition_id": result.condition_id,
                    "text_id": result.text_id,
                    "raw_output": result.raw_output,
                    "metadata": result.metadata,
                })

                out_path = run_dir / "raw" / cond_id / f"{text.id}.json"
                out_path.parent.mkdir(parents=True, exist_ok=True)
                out_path.write_text(json.dumps({
                    "condition": cond_id,
                    "text_id": text.id,
                    "raw_output": result.raw_output,
                    "metadata": result.metadata,
                }, indent=2, ensure_ascii=False), encoding="utf-8")
        return outputs

    def _load_cached_outputs(self, run_dir: Path) -> Dict[str, Dict[str, Any]]:
        from .conditions.base import ConditionOutput
        outputs: Dict[str, Dict[str, Any]] = {}
        raw_dir = run_dir / "raw"
        for cond_dir in raw_dir.iterdir():
            if not cond_dir.is_dir():
                continue
            cond_id = cond_dir.name
            for fp in cond_dir.glob("*.json"):
                data = json.loads(fp.read_text(encoding="utf-8"))
                tid = data["text_id"]
                outputs.setdefault(tid, {})[cond_id] = ConditionOutput(
                    condition_id=cond_id,
                    text_id=tid,
                    raw_output=data["raw_output"],
                    metadata=data.get("metadata", {}),
                )
        return outputs

    # ── Phase 3 helpers ───────────────────────────────────────────

    async def _evaluate_all(
        self,
        corpus: List[TestText],
        outputs: Dict[str, Dict[str, Any]],
        run_dir: Path,
    ) -> Dict[str, Dict[str, Dict]]:
        scores: Dict[str, Dict[str, Dict]] = {}
        for text in corpus:
            scores[text.id] = {}
            for cond_id in CONDITIONS:
                raw = sanitize_for_judge(
                    outputs[text.id][cond_id].raw_output, cond_id,
                )

                l1 = evaluate_layer1(raw, text.text)
                l2 = await evaluate_layer2(
                    self.judge, raw, text.text, l1.per_lens_detected,
                )
                ads = sum(l2.values())
                clis = await evaluate_layer3(self.judge, raw, text.text)
                l4 = await evaluate_layer4(self.judge, raw, text.text)

                scores[text.id][cond_id] = {
                    "lens_coverage": l1.lens_coverage,
                    "section_structure": l1.section_structure,
                    "term_precision": l1.term_precision,
                    "citation_density": l1.citation_density,
                    "length_ratio": l1.length_ratio,
                    "per_lens_depth": l2,
                    "ads": ads,
                    "clis": clis,
                    "discovery": l4.discovery_count,
                    "error_prevention": l4.error_prevention,
                }

                score_path = run_dir / "scores" / cond_id / f"{text.id}.json"
                score_path.parent.mkdir(parents=True, exist_ok=True)
                score_path.write_text(json.dumps(
                    scores[text.id][cond_id], indent=2, ensure_ascii=False,
                ), encoding="utf-8")
        return scores

    def _load_cached_scores(self, run_dir: Path) -> Dict[str, Dict[str, Dict]]:
        scores: Dict[str, Dict[str, Dict]] = {}
        for cond_dir in (run_dir / "scores").iterdir():
            if not cond_dir.is_dir():
                continue
            cond = cond_dir.name
            for fp in cond_dir.glob("*.json"):
                data = json.loads(fp.read_text(encoding="utf-8"))
                scores.setdefault(fp.stem, {})[cond] = data
        return scores

    # ── Phase 4 helpers ───────────────────────────────────────────

    async def _pairwise_all(
        self,
        corpus: List[TestText],
        outputs: Dict[str, Dict[str, Any]],
        run_dir: Path,
    ) -> List[Any]:
        all_pairwise: List[Any] = []
        for text in corpus:
            cond_outputs = {
                cid: sanitize_for_judge(outputs[text.id][cid].raw_output, cid)
                for cid in CONDITIONS
            }
            results = await run_all_pairwise(
                self.judge, text.id, text.text, cond_outputs,
            )
            all_pairwise.extend(results)

        pw_path = run_dir / "pairwise" / "all_comparisons.json"
        pw_path.parent.mkdir(parents=True, exist_ok=True)
        pw_path.write_text(json.dumps(
            [vars(r) for r in all_pairwise], indent=2, ensure_ascii=False,
        ), encoding="utf-8")

        bt_path = run_dir / "pairwise" / "bt_records.json"
        bt_path.write_text(json.dumps(
            self._pairwise_to_records(all_pairwise), indent=2, ensure_ascii=False,
        ), encoding="utf-8")

        return all_pairwise

    def _load_cached_pairwise(self, run_dir: Path) -> List[Any]:
        from .judges.pairwise import PairwiseResult
        raw = json.loads(
            (run_dir / "pairwise" / "all_comparisons.json").read_text(encoding="utf-8")
        )
        return [PairwiseResult(**r) for r in raw]

    # ── aggregation & H7/H8 ──────────────────────────────────────

    def _aggregate_scores(
        self, scores: Dict[str, Dict[str, Dict]],
    ) -> Dict[str, Dict[str, List[float]]]:
        agg: Dict[str, Dict[str, List[float]]] = {c: {} for c in CONDITIONS}
        metrics = ["lens_coverage", "ads", "clis", "discovery", "error_prevention"]
        for cond in CONDITIONS:
            for metric in metrics:
                agg[cond][metric] = [
                    scores[tid][cond][metric] for tid in scores
                ]
        return agg

    def _test_h7(
        self, corpus: List[TestText], scores: Dict,
        cond: str = "SCAFFOLD_FULL",
    ) -> Dict:
        easy_deltas: List[float] = []
        hard_deltas: List[float] = []
        for text in corpus:
            delta = (scores[text.id][cond]["ads"]
                     - scores[text.id]["NAKED"]["ads"])
            if text.difficulty == "easy":
                easy_deltas.append(delta)
            elif text.difficulty == "hard":
                hard_deltas.append(delta)
        result: Dict[str, Any] = {
            "mean_easy_delta": float(np.mean(easy_deltas)) if easy_deltas else 0,
            "mean_hard_delta": float(np.mean(hard_deltas)) if hard_deltas else 0,
        }
        if len(hard_deltas) > 5 and len(easy_deltas) > 5:
            result["test"] = wilcoxon_paired(hard_deltas, easy_deltas)
            result["effect"] = cliffs_delta(hard_deltas, easy_deltas)
        else:
            result["note"] = "insufficient_data"
        return result

    def _test_h8(
        self, corpus: List[TestText], scores: Dict,
        cond: str = "SCAFFOLD_FULL",
    ) -> Dict:
        domains = sorted({t.domain for t in corpus})
        per_domain: Dict[str, Dict[str, float]] = {}
        for domain in domains:
            domain_texts = [t for t in corpus if t.domain == domain]
            full = [scores[t.id][cond]["ads"] for t in domain_texts]
            naked = [scores[t.id]["NAKED"]["ads"] for t in domain_texts]
            per_domain[domain] = {
                "mean_full": float(np.mean(full)),
                "mean_naked": float(np.mean(naked)),
                "delta": float(np.mean(np.array(full) - np.array(naked))),
            }
        return per_domain

    @staticmethod
    def _pairwise_to_records(pairwise: List[Any]) -> List[Dict]:
        records: List[Dict] = []
        for r in pairwise:
            if r.true_winner == "TIE":
                records.append({
                    "winner": "TIE",
                    "cond_a": r.condition_a,
                    "cond_b": r.condition_b,
                })
            else:
                loser = (r.condition_b
                         if r.true_winner == r.condition_a
                         else r.condition_a)
                records.append({
                    "winner": r.true_winner,
                    "loser": loser,
                    "cond_a": r.condition_a,
                    "cond_b": r.condition_b,
                })
        return records
