# src/svy/core/adjustment_history.py
"""
Track history of weight adjustments for reproducibility and debugging.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any


@dataclass
class AdjustmentRecord:
    """
    Record of a single weight adjustment operation.

    Attributes:
        timestamp: When the adjustment was applied
        method: Type of adjustment (e.g., "nonresponse", "poststratify", "rake")
        params: Parameters used for the adjustment
        weight_col: Name of weight column that was adjusted
        rep_prefix: Prefix of replicate weights (if adjusted)
        n_reps: Number of replicates (if adjusted)
    """

    timestamp: datetime
    method: str
    params: dict[str, Any]
    weight_col: str
    rep_prefix: str | None = None
    n_reps: int | None = None

    def __str__(self) -> str:
        """Readable summary of the adjustment."""
        time_str = self.timestamp.strftime("%Y-%m-%d %H:%M:%S")
        base = f"[{time_str}] {self.method} → {self.weight_col}"

        if self.rep_prefix:
            base += f" + {self.n_reps} replicates ({self.rep_prefix}*)"

        # Show key parameters
        if self.params:
            key_params = []
            for k, v in self.params.items():
                if k in ("by", "controls", "factors", "margins"):
                    if isinstance(v, dict):
                        key_params.append(f"{k}={list(v.keys())}")
                    elif v is not None:
                        key_params.append(f"{k}={v}")
                elif k in ("resp_status", "unknown_to_inelig", "tol", "max_iter"):
                    if v is not None:
                        key_params.append(f"{k}={v}")

            if key_params:
                base += f" ({', '.join(key_params)})"

        return base

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "timestamp": self.timestamp.isoformat(),
            "method": self.method,
            "params": self.params,
            "weight_col": self.weight_col,
            "rep_prefix": self.rep_prefix,
            "n_reps": self.n_reps,
        }


@dataclass
class AdjustmentHistory:
    """
    Track all adjustments applied to sample weights.

    Attributes:
        records: List of adjustment records in chronological order
    """

    records: list[AdjustmentRecord] = field(default_factory=list)

    def add(
        self,
        method: str,
        params: dict[str, Any],
        weight_col: str,
        rep_prefix: str | None = None,
        n_reps: int | None = None,
    ) -> None:
        """Add a new adjustment record."""
        record = AdjustmentRecord(
            timestamp=datetime.now(),
            method=method,
            params=params.copy(),  # Copy to avoid mutation
            weight_col=weight_col,
            rep_prefix=rep_prefix,
            n_reps=n_reps,
        )
        self.records.append(record)

    def __str__(self) -> str:
        """Readable summary of all adjustments."""
        if not self.records:
            return "No adjustments applied"

        lines = [f"Adjustment History ({len(self.records)} operations):"]
        lines.extend(f"  {i + 1}. {record}" for i, record in enumerate(self.records))
        return "\n".join(lines)

    def __repr__(self) -> str:
        return f"AdjustmentHistory({len(self.records)} records)"

    def to_list(self) -> list[dict[str, Any]]:
        """Convert to list of dictionaries for serialization."""
        return [record.to_dict() for record in self.records]

    def summary(self) -> dict[str, Any]:
        """Get summary statistics of adjustments."""
        if not self.records:
            return {"total": 0}

        methods = {}
        for record in self.records:
            methods[record.method] = methods.get(record.method, 0) + 1

        return {
            "total": len(self.records),
            "methods": methods,
            "first": self.records[0].timestamp.isoformat(),
            "last": self.records[-1].timestamp.isoformat(),
            "final_weight": self.records[-1].weight_col,
        }
