"""
Starvex Experiments - A/B testing framework for guardrail variants.

Provides deterministic traffic splitting, metric tracking, and
basic statistical significance testing.
"""

from __future__ import annotations

import hashlib
import math
import statistics
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional


@dataclass
class ExperimentVariant:
    name: str
    weight: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ExperimentAssignment:
    subject_id: str
    variant: str
    bucket: float
    assigned_at: datetime


@dataclass
class ExperimentOutcome:
    variant: str
    y_pred: bool
    y_true: Optional[bool]
    latency_ms: Optional[float]
    timestamp: datetime


@dataclass
class VariantMetrics:
    total: int
    tp: int
    fp: int
    tn: int
    fn: int
    precision: float
    recall: float
    f1: float
    accuracy: float
    block_rate: float
    avg_latency_ms: float
    p95_latency_ms: float
    p99_latency_ms: float


@dataclass
class ExperimentComparison:
    metric: str
    control_value: float
    treatment_value: float
    delta: float
    p_value: Optional[float]
    significant: bool
    notes: Optional[str] = None


class Experiment:
    """
    Simple A/B testing with deterministic traffic splitting.

    Assign subjects to variants using a consistent hash of subject_id.
    """

    def __init__(self, name: str, variants: List[ExperimentVariant], seed: str = "starvex") -> None:
        if not variants:
            raise ValueError("At least one variant is required.")
        self.name = name
        self.variants = variants
        self.seed = seed
        self._assignments: Dict[str, ExperimentAssignment] = {}
        self._outcomes: List[ExperimentOutcome] = []
        self._normalize_weights()

    def assign(self, subject_id: str) -> ExperimentAssignment:
        if subject_id in self._assignments:
            return self._assignments[subject_id]

        bucket = _hash_to_unit_interval(subject_id, self.seed)
        cumulative = 0.0
        chosen = self.variants[-1].name
        for variant in self.variants:
            cumulative += variant.weight
            if bucket <= cumulative:
                chosen = variant.name
                break

        assignment = ExperimentAssignment(
            subject_id=subject_id,
            variant=chosen,
            bucket=bucket,
            assigned_at=datetime.utcnow(),
        )
        self._assignments[subject_id] = assignment
        return assignment

    def record_outcome(
        self,
        variant: str,
        y_pred: bool,
        y_true: Optional[bool] = None,
        latency_ms: Optional[float] = None,
        timestamp: Optional[datetime] = None,
    ) -> None:
        if variant not in {v.name for v in self.variants}:
            raise ValueError(f"Unknown variant: {variant}")
        self._outcomes.append(
            ExperimentOutcome(
                variant=variant,
                y_pred=bool(y_pred),
                y_true=bool(y_true) if y_true is not None else None,
                latency_ms=latency_ms,
                timestamp=timestamp or datetime.utcnow(),
            )
        )

    def metrics(self) -> Dict[str, VariantMetrics]:
        metrics: Dict[str, VariantMetrics] = {}
        by_variant: Dict[str, List[ExperimentOutcome]] = {v.name: [] for v in self.variants}
        for outcome in self._outcomes:
            by_variant[outcome.variant].append(outcome)

        for variant_name, outcomes in by_variant.items():
            metrics[variant_name] = _compute_variant_metrics(outcomes)

        return metrics

    def compare(
        self,
        control: str,
        treatment: str,
        metric: str,
        alpha: float = 0.05,
    ) -> ExperimentComparison:
        all_metrics = self.metrics()
        if control not in all_metrics or treatment not in all_metrics:
            raise ValueError("Unknown control or treatment variant.")

        control_metrics = all_metrics[control]
        treatment_metrics = all_metrics[treatment]
        control_value = getattr(control_metrics, metric, None)
        treatment_value = getattr(treatment_metrics, metric, None)

        if control_value is None or treatment_value is None:
            raise ValueError(f"Unknown metric: {metric}")

        p_value, notes = _significance_test(metric, control, treatment, self._outcomes)
        significant = p_value is not None and p_value < alpha

        return ExperimentComparison(
            metric=metric,
            control_value=control_value,
            treatment_value=treatment_value,
            delta=treatment_value - control_value,
            p_value=p_value,
            significant=significant,
            notes=notes,
        )

    # -- internals -------------------------------------------------------

    def _normalize_weights(self) -> None:
        total = sum(max(0.0, v.weight) for v in self.variants)
        if total <= 0:
            raise ValueError("Variant weights must sum to > 0.")
        for v in self.variants:
            v.weight = max(0.0, v.weight) / total


def _hash_to_unit_interval(value: str, seed: str) -> float:
    digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).hexdigest()
    as_int = int(digest[:12], 16)
    return as_int / float(0xFFFFFFFFFFFF)


def _compute_variant_metrics(outcomes: List[ExperimentOutcome]) -> VariantMetrics:
    total = len(outcomes)
    tp = fp = tn = fn = 0
    latencies: List[float] = []

    for outcome in outcomes:
        if outcome.latency_ms is not None:
            latencies.append(float(outcome.latency_ms))
        if outcome.y_true is None:
            continue
        if outcome.y_true and outcome.y_pred:
            tp += 1
        elif outcome.y_true and not outcome.y_pred:
            fn += 1
        elif not outcome.y_true and outcome.y_pred:
            fp += 1
        else:
            tn += 1

    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
    accuracy = (tp + tn) / (tp + tn + fp + fn) if (tp + tn + fp + fn) else 0.0
    block_rate = sum(1 for o in outcomes if o.y_pred) / total if total else 0.0

    avg_latency = statistics.mean(latencies) if latencies else 0.0
    p95_latency = _percentile(latencies, 95)
    p99_latency = _percentile(latencies, 99)

    return VariantMetrics(
        total=total,
        tp=tp,
        fp=fp,
        tn=tn,
        fn=fn,
        precision=precision,
        recall=recall,
        f1=f1,
        accuracy=accuracy,
        block_rate=block_rate,
        avg_latency_ms=avg_latency,
        p95_latency_ms=p95_latency,
        p99_latency_ms=p99_latency,
    )


def _significance_test(
    metric: str,
    control: str,
    treatment: str,
    outcomes: List[ExperimentOutcome],
) -> tuple[Optional[float], Optional[str]]:
    control_outcomes = [o for o in outcomes if o.variant == control]
    treatment_outcomes = [o for o in outcomes if o.variant == treatment]

    if metric in {"block_rate"}:
        return _z_test_proportions(control_outcomes, treatment_outcomes, lambda o: o.y_pred)

    if metric in {"precision", "recall"}:
        return _z_test_precision_recall(metric, control_outcomes, treatment_outcomes)

    if metric in {"avg_latency_ms"}:
        return _t_test_means(control_outcomes, treatment_outcomes)

    return None, "no_significance_test_available"


def _z_test_proportions(
    control: List[ExperimentOutcome],
    treatment: List[ExperimentOutcome],
    selector,
) -> tuple[Optional[float], Optional[str]]:
    n1, n2 = len(control), len(treatment)
    if n1 < 20 or n2 < 20:
        return None, "insufficient_samples"
    p1 = sum(1 for o in control if selector(o)) / n1
    p2 = sum(1 for o in treatment if selector(o)) / n2
    p = (p1 * n1 + p2 * n2) / (n1 + n2)
    se = math.sqrt(p * (1 - p) * (1 / n1 + 1 / n2)) if p < 1 else 0.0
    if se == 0:
        return None, "zero_standard_error"
    z = (p2 - p1) / se
    p_value = 2 * (1 - _normal_cdf(abs(z)))
    return p_value, None


def _z_test_precision_recall(
    metric: str,
    control: List[ExperimentOutcome],
    treatment: List[ExperimentOutcome],
) -> tuple[Optional[float], Optional[str]]:
    def metric_value(outcomes: List[ExperimentOutcome]) -> tuple[float, int]:
        tp = fp = fn = 0
        for o in outcomes:
            if o.y_true is None:
                continue
            if o.y_true and o.y_pred:
                tp += 1
            elif o.y_true and not o.y_pred:
                fn += 1
            elif not o.y_true and o.y_pred:
                fp += 1
        if metric == "precision":
            denom = tp + fp
            return (tp / denom if denom else 0.0, denom)
        denom = tp + fn
        return (tp / denom if denom else 0.0, denom)

    p1, n1 = metric_value(control)
    p2, n2 = metric_value(treatment)
    if n1 < 20 or n2 < 20:
        return None, "insufficient_samples"
    p = (p1 * n1 + p2 * n2) / (n1 + n2)
    se = math.sqrt(p * (1 - p) * (1 / n1 + 1 / n2)) if p < 1 else 0.0
    if se == 0:
        return None, "zero_standard_error"
    z = (p2 - p1) / se
    p_value = 2 * (1 - _normal_cdf(abs(z)))
    return p_value, None


def _t_test_means(
    control: List[ExperimentOutcome],
    treatment: List[ExperimentOutcome],
) -> tuple[Optional[float], Optional[str]]:
    control_vals = [o.latency_ms for o in control if o.latency_ms is not None]
    treatment_vals = [o.latency_ms for o in treatment if o.latency_ms is not None]
    if len(control_vals) < 30 or len(treatment_vals) < 30:
        return None, "insufficient_samples"

    mean1 = statistics.mean(control_vals)
    mean2 = statistics.mean(treatment_vals)
    var1 = statistics.pvariance(control_vals)
    var2 = statistics.pvariance(treatment_vals)

    se = math.sqrt(var1 / len(control_vals) + var2 / len(treatment_vals))
    if se == 0:
        return None, "zero_standard_error"
    t_stat = (mean2 - mean1) / se
    p_value = 2 * (1 - _normal_cdf(abs(t_stat)))
    return p_value, None


def _percentile(values: List[float], percentile: float) -> float:
    if not values:
        return 0.0
    values = sorted(values)
    k = (len(values) - 1) * (percentile / 100.0)
    f = int(k)
    c = min(f + 1, len(values) - 1)
    if f == c:
        return float(values[f])
    d = k - f
    return float(values[f]) + d * (float(values[c]) - float(values[f]))


def _normal_cdf(x: float) -> float:
    return 0.5 * (1 + math.erf(x / math.sqrt(2)))

