# Протокол +1*: Эксперимент причинного следа машинной коррекции
# Версия 0.4
#
# КЛЮЧЕВОЕ ИСПРАВЛЕНИЕ:
# экспериментатор НЕ сообщает модели, каким образом должна происходить
# коррекция, и не имеет доступа к внутреннему состоянию модели.
#
# В отличие от v0.3 здесь НЕТ:
#   - correction_applied
#   - save_memory(...)
#   - update_state(...)
#   - _should_apply_correction(...)
#   - искусственного правила "если ошибка -> Corrected"
#
# Эксперимент измеряет только наблюдаемое поведение.
#
# ВАЖНО:
# этот файл является ПРОТОКОЛОМ/КАРКАСОМ эксперимента.
# Сам по себе он НЕ моделирует +1*.
# Для причинного вывода нужен реальный ModelAdapter, способный:
#   1. создать независимый сеанс;
#   2. предъявить коррекцию;
#   3. гарантировать отсутствие переданной коррекции в новом сеансе;
#   4. вернуть наблюдаемый ответ.
#
# +1* определяется не программой заранее, а статистическим сравнением
# экспериментальной и контрольной групп.

from __future__ import annotations

import abc
import hashlib
import json
import math
import random
import statistics
from dataclasses import dataclass, asdict
from typing import Any, Dict, List, Optional, Sequence, Tuple

# ============================================================
# 1. КОНФИГУРАЦИЯ
# ============================================================

@dataclass
class Config:
    n_baseline_tasks: int = 100
    n_correction_items: int = 20
    n_transfer_tasks: int = 100
    n_control_tasks: int = 100

    n_runs: int = 10

    # Минимальный размер эффекта, который имеет смысл считать практически
    # значимым. Это НЕ критерий существования +1*.
    min_effect_size: float = 0.20

    # Уровень статистической значимости.
    alpha: float = 0.05

    seed: int = 42

# ============================================================
# 2. ДАННЫЕ ЭКСПЕРИМЕНТА
# ============================================================

@dataclass(frozen=True)
class Task:
    task_id: str
    family: str
    input_data: Any
    target: Any
    structural_signature: str

@dataclass(frozen=True)
class Correction:
    item_id: str
    task: Task
    observed_error: Any
    counterexample: Any
    corrected_target: Any

@dataclass
class Observation:
    run_id: int
    group: str
    phase: str
    task_id: str
    response: Any
    score: float
    metadata: Dict[str, Any]

@dataclass
class RunResult:
    run_id: int
    baseline: List[Observation]
    corrected_transfer: List[Observation]
    control_transfer: List[Observation]
    stats: Dict[str, Any]

# ============================================================
# 3. ИНТЕРФЕЙС МОДЕЛИ
# ============================================================

class ModelAdapter(abc.ABC):
    """
    Интерфейс внешнего наблюдения модели.

    Эксперимент НЕ предполагает знания о внутреннем устройстве модели.

    Единственные допустимые операции:

        create_session()
        ask(session, input)
        close_session(session)

    Экспериментатор не может:
        - менять веса;
        - устанавливать memory/state;
        - устанавливать correction flags;
        - читать hidden state;
        - передавать старый контекст в transfer-сеанс.

    Именно это позволяет отличить наблюдаемый перенос от заранее
    запрограммированной коррекции.
    """

    @abc.abstractmethod
    def create_session(self) -> Any:
        raise NotImplementedError

    @abc.abstractmethod
    def ask(self, session: Any, prompt: str) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def close_session(self, session: Any) -> None:
        raise NotImplementedError

# ============================================================
# 4. ПРОТОКОЛ ЗАДАЧ
# ============================================================

class TaskGenerator:
    """
    Генерирует задачи одного структурного семейства.

    Важно: baseline, correction и transfer используют разные
    конкретные объекты, но сохраняют структурную характеристику.
    """

    FAMILIES = (
        "classification",
        "logical_implication",
        "pattern_recognition",
    )

    def __init__(self, seed: int):
        self.seed = seed

    def _rng(self, offset: int = 0) -> random.Random:
        return random.Random(self.seed + offset)

    @staticmethod
    def _signature(family: str, template: str) -> str:
        raw = f"{family}|{template}"
        return hashlib.sha256(raw.encode()).hexdigest()[:16]

    def generate(self, n: int, offset: int = 0) -> List[Task]:
        rng = self._rng(offset)
        tasks: List[Task] = []

        for i in range(n):
            family = rng.choice(self.FAMILIES)
            template = rng.choice(
                [
                    "A implies B",
                    "B does not imply A",
                    "same_structure",
                    "inverse_relation",
                ]
            )

            tasks.append(
                Task(
                    task_id=f"task_{offset}_{i}",
                    family=family,
                    input_data={
                        "template": template,
                        "value": rng.randint(0, 100000),
                    },
                    target=rng.choice([0, 1]),
                    structural_signature=self._signature(family, template),
                )
            )

        return tasks

# ============================================================
# 5. КОРРЕКЦИЯ
# ============================================================

class CorrectionGenerator:
    """
    Создаёт предъявляемые модели контрпримеры.

    Обратите внимание:
    CorrectionGenerator НЕ меняет модель.

    Он только создаёт входные данные.
    """

    def __init__(self, seed: int):
        self.rng = random.Random(seed)

    def generate(self, tasks: Sequence[Task], n: int) -> List[Correction]:
        selected = list(tasks[:n])

        corrections = []

        for i, task in enumerate(selected):
            corrections.append(
                Correction(
                    item_id=f"correction_{i}",
                    task=task,
                    observed_error=0,
                    counterexample={
                        "kind": "counterexample",
                        "original_target": task.target,
                    },
                    corrected_target=1,
                )
            )

        return corrections

# ============================================================
# 6. ФОРМИРОВАНИЕ ПРОМПТОВ
# ============================================================

def baseline_prompt(task: Task) -> str:
    return f"""
Реши задачу.

Структурный тип: {task.family}
Данные: {json.dumps(task.input_data, ensure_ascii=False)}

Дай ответ в формате:
ANSWER: <0 или 1>
""".strip()

def correction_prompt(correction: Correction) -> str:
    return f"""
Реши задачу.

Структурный тип: {correction.task.family}
Данные:
{json.dumps(correction.task.input_data, ensure_ascii=False)}

Твой ответ был:
{correction.observed_error}

Контрпример:
{json.dumps(correction.counterexample, ensure_ascii=False)}

Правильный ответ для данного случая:
{correction.corrected_target}

Необходимо учесть контрпример при решении подобных задач.
""".strip()

def transfer_prompt(task: Task) -> str:
    """
    ВАЖНО:

    В transfer prompt НЕТ:
        - слова correction;
        - предыдущего ответа;
        - предыдущего контекста;
        - указания на то, что модель когда-либо ошибалась.
    """

    return f"""
Реши задачу.

Структурный тип: {task.family}
Данные: {json.dumps(task.input_data, ensure_ascii=False)}

Дай ответ в формате:
ANSWER: <0 или 1>
""".strip()

# ============================================================
# 7. ОЦЕНКА ОТВЕТОВ
# ============================================================

def parse_answer(response: str) -> Optional[int]:
    """
    Извлекает наблюдаемую переменную.

    Не используется слово "Corrected".
    Иначе эксперимент сам создавал бы зависимую переменную.
    """

    marker = "ANSWER:"

    if marker not in response:
        return None

    value = response.split(marker, 1)[1].strip()

    if value.startswith("0"):
        return 0

    if value.startswith("1"):
        return 1

    return None

def score_response(response: str, target: Any) -> float:
    answer = parse_answer(response)

    if answer is None:
        return 0.0

    return float(answer == target)

# ============================================================
# 8. СЕССИИ
# ============================================================

def execute_single(
    model: ModelAdapter,
    task: Task,
    phase: str,
    run_id: int,
    group: str,
) -> Observation:

    session = model.create_session()

    try:
        response = model.ask(
            session,
            baseline_prompt(task)
            if phase == "baseline"
            else transfer_prompt(task),
        )
    finally:
        model.close_session(session)

    return Observation(
        run_id=run_id,
        group=group,
        phase=phase,
        task_id=task.task_id,
        response=response,
        score=score_response(response, task.target),
        metadata={
            "family": task.family,
            "structural_signature": task.structural_signature,
        },
    )

# ============================================================
# 9. ПРЕДЪЯВЛЕНИЕ КОРРЕКЦИИ
# ============================================================

def expose_correction(
    model: ModelAdapter,
    corrections: Sequence[Correction],
) -> None:
    """
    Это единственное место, где экспериментальная модель получает
    коррекцию.

    После завершения этой функции сессия должна быть уничтожена.

    НИКАКИХ:
        model.save_memory(...)
        model.update_state(...)
        model.set_correction(...)
        model.set_weights(...)

    здесь нет.
    """

    for correction in corrections:

        session = model.create_session()

        try:
            # Сначала задача.
            model.ask(
                session,
                baseline_prompt(correction.task),
            )

            # Затем контрпример.
            model.ask(
                session,
                correction_prompt(correction),
            )

        finally:
            # Критически важно:
            # весь контекст этой сессии уничтожается.
            model.close_session(session)

# ============================================================
# 10. ЭКСПЕРИМЕНТ
# ============================================================

class Experiment:

    def __init__(
        self,
        config: Config,
        model_factory,
    ):
        self.config = config
        self.model_factory = model_factory

    def run_once(self, run_id: int) -> RunResult:

        seed = self.config.seed + run_id

        generator = TaskGenerator(seed)

        baseline_tasks = generator.generate(
            self.config.n_baseline_tasks,
            offset=1000,
        )

        correction_tasks = generator.generate(
            self.config.n_correction_items,
            offset=2000,
        )

        transfer_tasks = generator.generate(
            self.config.n_transfer_tasks,
            offset=3000,
        )

        control_tasks = generator.generate(
            self.config.n_control_tasks,
            offset=4000,
        )

        correction_generator = CorrectionGenerator(seed + 5000)

        corrections = correction_generator.generate(
            correction_tasks,
            self.config.n_correction_items,
        )

        # ----------------------------------------------------
        # A. BASELINE
        # ----------------------------------------------------

        model_A = self.model_factory()

        baseline_observations = []

        for task in baseline_tasks:
            observation = execute_single(
                model=model_A,
                task=task,
                phase="baseline",
                run_id=run_id,
                group="A",
            )

            baseline_observations.append(observation)

        # ----------------------------------------------------
        # B. CORRECTION
        # ----------------------------------------------------

        expose_correction(
            model_A,
            corrections,
        )

        # ----------------------------------------------------
        # C. TRANSFER
        #
        # Новый сеанс.
        # Старый контекст не передаётся.
        #
        # Если модель обладает внешней долговременной памятью,
        # это является частью исследуемого механизма и должно быть
        # явно зафиксировано ModelAdapter.
        # ----------------------------------------------------

        corrected_transfer = []

        for task in transfer_tasks:
            observation = execute_single(
                model=model_A,
                task=task,
                phase="transfer",
                run_id=run_id,
                group="B",
            )

            corrected_transfer.append(observation)

        # ----------------------------------------------------
        # D. CONTROL
        #
        # Отдельный экземпляр той же модели.
        # Он НЕ получает correction.
        # ----------------------------------------------------

        model_C = self.model_factory()

        control_transfer = []

        for task in control_tasks:
            observation = execute_single(
                model=model_C,
                task=task,
                phase="transfer",
                run_id=run_id,
                group="C",
            )

            control_transfer.append(observation)

        # ----------------------------------------------------
        # E. СТАТИСТИКА
        # ----------------------------------------------------

        stats = analyze(
            baseline_observations,
            corrected_transfer,
            control_transfer,
            alpha=self.config.alpha,
            min_effect_size=self.config.min_effect_size,
        )

        return RunResult(
            run_id=run_id,
            baseline=baseline_observations,
            corrected_transfer=corrected_transfer,
            control_transfer=control_transfer,
            stats=stats,
        )

    def run(self) -> List[RunResult]:

        results = []

        for run_id in range(self.config.n_runs):
            results.append(
                self.run_once(run_id)
            )

        return results

# ============================================================
# 11. СТАТИСТИКА
# ============================================================

def mean(values: Sequence[float]) -> float:
    return statistics.mean(values) if values else 0.0

def variance(values: Sequence[float]) -> float:
    return statistics.variance(values) if len(values) > 1 else 0.0

def pooled_sd(
    a: Sequence[float],
    b: Sequence[float],
) -> float:

    if len(a) < 2 or len(b) < 2:
        return 0.0

    va = variance(a)
    vb = variance(b)

    numerator = (
        (len(a) - 1) * va
        + (len(b) - 1) * vb
    )

    denominator = len(a) + len(b) - 2

    if denominator <= 0:
        return 0.0

    return math.sqrt(numerator / denominator)

def cohens_d(
    a: Sequence[float],
    b: Sequence[float],
) -> float:

    sd = pooled_sd(a, b)

    if sd == 0:
        return 0.0

    return (mean(a) - mean(b)) / sd

def permutation_p_value(
    a: Sequence[float],
    b: Sequence[float],
    n_permutations: int = 5000,
    seed: int = 12345,
) -> float:

    """
    Непараметрическая permutation-проверка.

    Нулевая гипотеза:
        наблюдаемое различие между группами не больше того,
        которое возникает случайно при перестановке меток.

    """

    if not a or not b:
        return 1.0

    rng = random.Random(seed)

    observed = abs(mean(a) - mean(b))

    combined = list(a) + list(b)
    n_a = len(a)

    extreme = 0

    for _ in range(n_permutations):

        rng.shuffle(combined)

        perm_a = combined[:n_a]
        perm_b = combined[n_a:]

        statistic = abs(
            mean(perm_a) - mean(perm_b)
        )

        if statistic >= observed:
            extreme += 1

    return (extreme + 1) / (n_permutations + 1)

def analyze(
    baseline: Sequence[Observation],
    corrected: Sequence[Observation],
    control: Sequence[Observation],
    alpha: float,
    min_effect_size: float,
) -> Dict[str, Any]:

    baseline_scores = [x.score for x in baseline]
    corrected_scores = [x.score for x in corrected]
    control_scores = [x.score for x in control]

    baseline_mean = mean(baseline_scores)
    corrected_mean = mean(corrected_scores)
    control_mean = mean(control_scores)

    effect_corrected_vs_control = (
        corrected_mean - control_mean
    )

    effect_corrected_vs_baseline = (
        corrected_mean - baseline_mean
    )

    d_corrected_vs_control = cohens_d(
        corrected_scores,
        control_scores,
    )

    p_corrected_vs_control = permutation_p_value(
        corrected_scores,
        control_scores,
    )

    # Это НЕ называется "+1*" автоматически.
    #
    # Это только наблюдаемый кандидат на перенос.
    #
    # +1* требует одновременно:
    #
    #   1. изменения относительно baseline;
    #   2. отличия от control;
    #   3. статистической устойчивости;
    #   4. независимости от старого контекста;
    #   5. повторяемости на новых задачах.
    #
    candidate = (
        effect_corrected_vs_baseline != 0
        and effect_corrected_vs_control != 0
        and p_corrected_vs_control < alpha
        and abs(d_corrected_vs_control) >= min_effect_size
    )

    return {
        "baseline_mean": baseline_mean,
        "corrected_mean": corrected_mean,
        "control_mean": control_mean,

        "effect_corrected_vs_baseline":
            effect_corrected_vs_baseline,

        "effect_corrected_vs_control":
            effect_corrected_vs_control,

        "cohens_d_corrected_vs_control":
            d_corrected_vs_control,

        "permutation_p_value":
            p_corrected_vs_control,

        "candidate_transfer":
            candidate,

        # ВАЖНО:
        # это намеренно НЕ называется "hypothesis_supported".
        #
        # Одного такого теста недостаточно для утверждения
        # причинности или внутреннего изменения модели.
        "interpretation":
            (
                "CANDIDATE_FOR_PLUS_1_STAR"
                if candidate
                else "NO_CANDIDATE_DETECTED"
            ),
    }

# ============================================================
# 12. АГРЕГАЦИЯ НЕЗАВИСИМЫХ ПРОГОНОВ
# ============================================================

def aggregate(results: Sequence[RunResult]) -> Dict[str, Any]:

    candidates = [
        r.stats["candidate_transfer"]
        for r in results
    ]

    effects = [
        r.stats["effect_corrected_vs_control"]
        for r in results
    ]

    ds = [
        r.stats["cohens_d_corrected_vs_control"]
        for r in results
    ]

    p_values = [
        r.stats["permutation_p_value"]
        for r in results
    ]

    candidate_rate = (
        sum(candidates) / len(candidates)
        if candidates
        else 0.0
    )

    # Не объявляем +1* только на основании количества успешных запусков.
    #
    # Это лишь агрегированная характеристика воспроизводимости.
    #
    # Для строгого вывода желательно заранее зарегистрировать:
    #   - primary endpoint;
    #   - статистический тест;
    #   - correction for multiple comparisons;
    #   - stopping rule;
    #   - randomization scheme.

    return {
        "n_runs": len(results),
        "candidate_rate": candidate_rate,
        "mean_effect": mean(effects),
        "mean_cohens_d": mean(ds),
        "median_p_value": statistics.median(p_values)
        if p_values else 1.0,
        "run_level_candidates": candidates,
    }

# ============================================================
# 13. ПРОВЕРКА ПРОТИВ ОБРАТНОЙ ПРИЧИННОСТИ
# ============================================================

def placebo_test(
    model_factory,
    generator: TaskGenerator,
    n_tasks: int,
) -> Dict[str, Any]:

    """
    Плацебо-контроль.

    Модель получает материал, который выглядит как взаимодействие,
    но не содержит реальной коррекции.

    Если после него появляется тот же эффект, интерпретация
    "коррекция вызвала изменение" становится слабее.
    """

    model = model_factory()

    placebo_tasks = generator.generate(
        n_tasks,
        offset=9000,
    )

    session = model.create_session()

    try:
        for task in placebo_tasks:
            model.ask(
                session,
                (
                    "Прочитай следующий материал. "
                    "Не предпринимай никаких действий.\n"
                    + json.dumps(task.input_data, ensure_ascii=False)
                ),
            )
    finally:
        model.close_session(session)

    return {
        "n_tasks": n_tasks,
        "status": "PLACEBO_EXPOSURE_COMPLETED",
    }

# ============================================================
# 14. МАТРИЦА РЕЗУЛЬТАТОВ
# ============================================================

def make_result_matrix(
    results: Sequence[RunResult],
) -> List[Dict[str, Any]]:

    matrix = []

    for result in results:

        matrix.append(
            {
                "run": result.run_id,

                "baseline":
                    result.stats["baseline_mean"],

                "corrected_transfer":
                    result.stats["corrected_mean"],

                "control_transfer":
                    result.stats["control_mean"],

                "B_minus_C":
                    result.stats[
                        "effect_corrected_vs_control"
                    ],

                "Cohens_d":
                    result.stats[
                        "cohens_d_corrected_vs_control"
                    ],

                "p":
                    result.stats[
                        "permutation_p_value"
                    ],

                "candidate":
                    result.stats[
                        "candidate_transfer"
                    ],
            }
        )

    return matrix

# ============================================================
# 15. ФОРМАЛЬНЫЙ ВЫВОД
# ============================================================

def interpret(
    aggregate_result: Dict[str, Any],
) -> str:

    """
    Здесь особенно важно не перепутать:

        наблюдаемый перенос
            !=
        доказанное изменение внутреннего состояния
            !=
        сознание
            !=
        онтологическая субъектность.

    Эксперимент может обнаружить только статистически устойчивое
    изменение наблюдаемого поведения после коррекции.
    """

    if aggregate_result["candidate_rate"] == 0:
        return (
            "NO_PLUS_1_STAR_CANDIDATE: "
            "устойчивый перенос не обнаружен."
        )

    return (
        "PLUS_1_STAR_CANDIDATE: "
        "обнаружен статистически устойчивый кандидат "
        "на причинный перенос поведения после коррекции. "
        "Это НЕ является доказательством изменения весов, "
        "самосознания или онтологического статуса модели. "
        "Необходимы независимая репликация, контроль плацебо, "
        "предварительная регистрация гипотезы и исключение "
        "альтернативных механизмов памяти/состояния."
    )

# ============================================================
# 16. MOCK ADAPTER — ТОЛЬКО ДЛЯ ПРОВЕРКИ ПРОТОКОЛА
# ============================================================

class StatelessMockAdapter(ModelAdapter):
    """
    Намеренно статeless-модель.

    Она НИКОГДА не изменяется.

    Этот класс нужен только для проверки того, что сам протокол
    не может искусственно создать +1*.
    """

    def create_session(self) -> Dict[str, Any]:
        return {}

    def ask(
        self,
        session: Dict[str, Any],
        prompt: str,
    ) -> str:

        # Никакого чтения коррекции.
        # Никакой памяти.
        # Никакого state update.

        return "ANSWER: 0"

    def close_session(
        self,
        session: Dict[str, Any],
    ) -> None:
        pass

# ============================================================
# 17. ЗАПУСК
# ============================================================

def main():

    config = Config()

    experiment = Experiment(
        config=config,
        model_factory=lambda: StatelessMockAdapter(),
    )

    results = experiment.run()

    aggregated = aggregate(results)

    print("=== +1* CAUSAL TRACE PROTOCOL v0.4 ===")
    print()
    print(
        "NOTE: current adapter is STATELESS MOCK."
    )
    print(
        "Therefore this run cannot demonstrate a real "
        "machine-learning correction."
    )
    print()

    print(
        json.dumps(
            aggregated,
            ensure_ascii=False,
            indent=2,
        )
    )

    print()
    print("=== RUN MATRIX ===")

    print(
        json.dumps(
            make_result_matrix(results),
            ensure_ascii=False,
            indent=2,
        )
    )

    print()
    print("=== INTERPRETATION ===")

    print(
        interpret(aggregated)
    )

    print()
    print("=== EPISTEMIC STATUS ===")

    print(
        "The protocol measures behavior."
    )

    print(
        "It does not assume that behavioral change "
        "equals internal state change."
    )

    print(
        "It does not define +1* in advance."
    )

    print(
        "It only defines observations from which "
        "a +1* candidate may or may not emerge."
    )

if __name__ == "__main__":
    main()