from __future__ import annotations

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize_scalar
from scipy.special import expit, logit
from typing import Tuple

__all__ = [
    "optimise_temperature",
    "calibrate",
    "reliability_curve",
    "EDGE_THR",
]

# ------------------------------------------------------------------
# Hyper‑parameters
# ------------------------------------------------------------------
# Minimum calibrated edge magnitude before we take a trade.
# EDGE_THR: float = 0.55  # Default value if not passed


# ------------------------------------------------------------------
#  Temperature scaling
# ------------------------------------------------------------------

def _nll_temperature(T: float, logit_p: np.ndarray, y_true: np.ndarray) -> float:
    """Negative log‑likelihood of labels given scaled logits."""
    p_cal = expit(logit_p / T)
    # Binary cross‑entropy (NLL)
    eps = 1e-12
    p_cal = np.clip(p_cal, eps, 1 - eps)
    nll = -(y_true * np.log(p_cal) + (1 - y_true) * np.log(1 - p_cal))
    return float(np.mean(nll))


def optimise_temperature(p_raw: np.ndarray, y_true: np.ndarray) -> float:
    """Return optimal temperature `T` that minimises NLL on `y_true`."""
    p_raw = p_raw.flatten().astype(float)
    y_true = y_true.flatten().astype(int)
    logit_p = logit(np.clip(p_raw, 1e-6, 1 - 1e-6))

    res = minimize_scalar(
        lambda T: _nll_temperature(T, logit_p, y_true),
        bounds=(0.05, 10.0),
        method="bounded",
    )
    return float(res.x)


# ------------------------------------------------------------------
#  Public API
# ------------------------------------------------------------------

def calibrate(p_raw: np.ndarray, T: float) -> np.ndarray:
    """Return temperature‑scaled probabilities."""
    return expit(logit(np.clip(p_raw, 1e-6, 1 - 1e-6)) / T)


def reliability_curve(
    p_raw: np.ndarray,
    y_true: np.ndarray,
    n_bins: int = 10,
    show: bool = False,
) -> Tuple[np.ndarray, np.ndarray]:
    """Return (bin_centres, empirical_prob) and optionally plot reliability."""
    p_raw = p_raw.flatten()
    y_true = y_true.flatten()
    bins = np.linspace(0, 1, n_bins + 1)
    bin_ids = np.digitize(p_raw, bins) - 1  # 0‑indexed

    bin_centres = 0.5 * (bins[:-1] + bins[1:])
    acc = np.zeros(n_bins)
    for i in range(n_bins):
        idx = bin_ids == i
        if np.any(idx):
            acc[i] = y_true[idx].mean()
        else:
            acc[i] = np.nan

    if show:
        plt.figure(figsize=(4, 4))
        plt.plot([0, 1], [0, 1], "k--", label="perfect")
        plt.plot(bin_centres, acc, "o-", label="empirical")
        plt.xlabel("Predicted P(up)")
        plt.ylabel("Actual frequency")
        plt.title("Reliability curve")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
    return bin_centres, acc


# ------------------------------------------------------------------
#  Signal filter
# ------------------------------------------------------------------

def action_signal(p_cal: np.ndarray, edge_threshold: float = 0.55) -> np.ndarray:
    """Return trading signal: 1, -1 or 0 based on calibrated edge threshold."""
    up = p_cal > edge_threshold
    dn = p_cal < 1 - edge_threshold
    return np.where(up, 1, np.where(dn, -1, 0)) 