"""Vector learner — learns model vectors from calibration results."""

from __future__ import annotations

import logging
from typing import Optional

import numpy as np

logger = logging.getLogger("ecate.calibration.vector_learner")


def learn_model_vector(
    task_embeddings: np.ndarray,
    pass_fail: np.ndarray,
    method: str = "logistic",
    regularization: float = 1.0,
) -> tuple[np.ndarray, float]:
    """Learn a model vector from task embeddings and pass/fail results.
    
    Args:
        task_embeddings: [n_tasks, embed_dim] task embeddings
        pass_fail: [n_tasks] binary array (1=pass, 0=fail)
        method: "logistic" or "linear"
        regularization: Regularization strength (higher = more regularization)
    
    Returns:
        (model_vector, accuracy) where model_vector has shape (embed_dim,)
    """
    if len(task_embeddings) != len(pass_fail):
        raise ValueError(f"Mismatched lengths: {len(task_embeddings)} vs {len(pass_fail)}")
    
    n_samples = len(pass_fail)
    n_pass = int(pass_fail.sum())
    n_fail = n_samples - n_pass
    
    logger.info(f"Learning model vector from {n_samples} samples ({n_pass} pass, {n_fail} fail)")
    
    if n_pass == 0 or n_fail == 0:
        logger.warning("All samples have same label, returning zero vector")
        return np.zeros(task_embeddings.shape[1]), 0.0
    
    if method == "logistic":
        return _learn_logistic(task_embeddings, pass_fail, regularization)
    else:
        return _learn_linear(task_embeddings, pass_fail, regularization)


def _learn_logistic(
    X: np.ndarray, 
    y: np.ndarray, 
    C: float = 1.0,
) -> tuple[np.ndarray, float]:
    """Learn using logistic regression."""
    try:
        from sklearn.linear_model import LogisticRegression
        from sklearn.preprocessing import StandardScaler
    except ImportError:
        logger.warning("sklearn not available, falling back to linear")
        return _learn_linear(X, y, C)
    
    # Scale features for better convergence
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    clf = LogisticRegression(
        C=C,
        max_iter=1000,
        class_weight="balanced",
        solver="lbfgs",
        random_state=42,
    )
    clf.fit(X_scaled, y)
    
    # Get accuracy
    accuracy = clf.score(X_scaled, y)
    
    # Transform coefficients back to original scale
    # coef_original = coef_scaled / scale
    coef = clf.coef_[0] / scaler.scale_
    
    # Normalize to unit length
    norm = np.linalg.norm(coef)
    if norm > 0:
        coef = coef / norm
    
    logger.info(f"Logistic regression: accuracy={accuracy:.3f}")
    return coef, accuracy


def _learn_linear(
    X: np.ndarray, 
    y: np.ndarray, 
    alpha: float = 1.0,
) -> tuple[np.ndarray, float]:
    """Learn using ridge regression (fallback if sklearn not available)."""
    # Ridge regression: (X.T @ X + αI)^-1 @ X.T @ y
    n_features = X.shape[1]
    
    # Center y
    y_centered = y - y.mean()
    
    # Solve with regularization
    XtX = X.T @ X
    XtX += alpha * np.eye(n_features)
    Xty = X.T @ y_centered
    
    try:
        coef = np.linalg.solve(XtX, Xty)
    except np.linalg.LinAlgError:
        # Fallback to pseudoinverse
        coef = np.linalg.lstsq(X, y_centered, rcond=None)[0]
    
    # Normalize
    norm = np.linalg.norm(coef)
    if norm > 0:
        coef = coef / norm
    
    # Calculate accuracy (threshold at 0.5)
    preds = (X @ coef + y.mean() > 0.5).astype(float)
    accuracy = (preds == y).mean()
    
    logger.info(f"Linear regression: accuracy={accuracy:.3f}")
    return coef, accuracy


def update_model_vector_incremental(
    existing_vector: np.ndarray,
    new_embeddings: np.ndarray,
    new_pass_fail: np.ndarray,
    learning_rate: float = 0.1,
) -> np.ndarray:
    """Update an existing model vector with new calibration data.
    
    Uses exponential moving average to blend old and new information.
    Useful for adding small batches of new data without full retraining.
    
    Args:
        existing_vector: Current model vector [embed_dim]
        new_embeddings: New task embeddings [n_new, embed_dim]
        new_pass_fail: Binary results [n_new]
        learning_rate: How much to weight new data (0-1)
    
    Returns:
        Updated model vector [embed_dim]
    """
    # Learn vector from new data
    new_vector, _ = learn_model_vector(new_embeddings, new_pass_fail)
    
    # Blend with existing
    updated = (1 - learning_rate) * existing_vector + learning_rate * new_vector
    
    # Normalize
    norm = np.linalg.norm(updated)
    if norm > 0:
        updated = updated / norm
    
    return updated
