#!/usr/bin/env python3
"""
linear_scaling.py — Keijzer Linear Scaling (2003)

Given expression g(x) and target y, solve for optimal scale and offset:
    f(x) = a·g(x) + b

Closed-form solution via least squares:
    a = cov(g, y) / var(g)
    b = mean(y) - a·mean(g)

This reduces the search space: GP only needs to find the SHAPE g(x),
not the exact scale and offset.
"""

from __future__ import annotations

import numpy as np
from typing import Tuple


def linear_scale(
    preds: np.ndarray,
    target: np.ndarray,
) -> Tuple[float, float, np.ndarray]:
    """Solve for optimal a, b in f(x) = a·g(x) + b.

    Args:
        preds: raw expression predictions g(x_i)
        target: target values y_i

    Returns:
        (a, b, scaled_preds) where scaled_preds = a·preds + b
    """
    # Handle degenerate cases
    if len(preds) < 2:
        return 1.0, 0.0, preds

    pred_std = np.std(preds)
    if pred_std < 1e-15:
        # Constant expression: best fit is mean(y)
        return 0.0, float(np.mean(target)), np.full_like(target, np.mean(target))

    # Least squares: a = cov(g,y)/var(g), b = mean(y) - a*mean(g)
    pred_mean = np.mean(preds)
    target_mean = np.mean(target)

    cov_gy = np.mean((preds - pred_mean) * (target - target_mean))
    var_g = np.mean((preds - pred_mean) ** 2)

    if var_g < 1e-15:
        return 0.0, target_mean, np.full_like(target, target_mean)

    a = cov_gy / var_g
    b = target_mean - a * pred_mean

    scaled = a * preds + b
    return float(a), float(b), scaled


def compute_r2(
    preds: np.ndarray,
    target: np.ndarray,
) -> float:
    """Compute R² (coefficient of determination).

    R² = 1 - SS_res / SS_tot
    """
    if len(target) < 2:
        return 0.0

    ss_res = np.sum((target - preds) ** 2)
    ss_tot = np.sum((target - np.mean(target)) ** 2)

    if ss_tot < 1e-15:
        return 1.0 if ss_res < 1e-15 else 0.0

    return float(1.0 - ss_res / ss_tot)


def compute_mse(
    preds: np.ndarray,
    target: np.ndarray,
) -> float:
    """Compute mean squared error."""
    return float(np.mean((target - preds) ** 2))