from __future__ import annotations

from typing import Optional, List, Union
from typing import Any, Dict, Optional, List

import numpy as np
import pandas as pd
from pydantic import BaseModel, ConfigDict


class DiagnosticData(BaseModel):
    """Base class for all diagnostic data_contracts."""

    model_config = ConfigDict(arbitrary_types_allowed=True)


class UnconfoundednessDiagnosticData(DiagnosticData):
    """Fields common to all models assuming unconfoundedness."""

    m_hat: np.ndarray  # Propensity scores
    d: np.ndarray  # Treatment indicators
    y: Optional[np.ndarray] = None  # Outcomes
    x: Optional[np.ndarray] = None  # Confounders (for balance checks)
    g0_hat: Optional[np.ndarray] = None  # Estimated outcome under control
    g1_hat: Optional[np.ndarray] = None  # Estimated outcome under treatment
    psi_b: Optional[np.ndarray] = None  # Orthogonal signal (for DML)
    folds: Optional[np.ndarray] = None  # Cross-fitting folds
    trimming_threshold: float = 0.0

    # Sensitivity elements (DoubleML-style)
    sigma2: Optional[float] = None
    nu2: Optional[float] = None
    psi_sigma2: Optional[np.ndarray] = None
    psi_nu2: Optional[np.ndarray] = None
    riesz_rep: Optional[np.ndarray] = None
    m_alpha: Optional[np.ndarray] = None
    psi: Optional[np.ndarray] = None
    score: Optional[str] = None  # ATE or ATTE
    sensitivity_analysis: Optional[Dict[str, Any]] = None


class MultiUnconfoundednessDiagnosticData(DiagnosticData):
    """Fields common to all models assuming unconfoundedness with multi_uncofoundedness."""

    m_hat: np.ndarray  # Propensity scores
    d: np.ndarray  # Treatments indicators
    y: Optional[np.ndarray] = None  # Outcomes
    x: Optional[np.ndarray] = None  # Confounders (for balance checks)
    g_hat: Optional[np.ndarray] = None  # Estimated outcome under control
    psi_b: Optional[np.ndarray] = None  # Orthogonal signal (for DML)
    folds: Optional[np.ndarray] = None  # Cross-fitting folds
    trimming_threshold: float = 0.0

    # Sensitivity elements (DoubleML-style)
    sigma2: Union[float, np.ndarray] = None
    nu2: Optional[np.ndarray] = None
    psi_sigma2: Optional[np.ndarray] = None
    psi_nu2: Optional[np.ndarray] = None
    riesz_rep: Optional[np.ndarray] = None
    m_alpha: Optional[np.ndarray] = None
    psi: Optional[np.ndarray] = None
    score: Optional[str] = None  # ATE or ATTE


class DiffInMeansDiagnosticData(DiagnosticData):
    """Diagnostic data_contracts for Difference-in-Means model."""

    pass


class CUPEDDiagnosticData(DiagnosticData):
    """Diagnostic data_contracts for CUPED / ANCOVA models."""

    ate_naive: float
    se_naive: float
    variance_reduction_pct: float
    beta_covariates: np.ndarray
    gamma_interactions: np.ndarray
    covariates: List[str]
    adj_type: str
