from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable

from scipy.optimize import minimize

from app.models import Component

ShareTuple = tuple[float, float, float]
ScoreFn = Callable[[ShareTuple], float]


@dataclass
class OptimizeSharesResult:
    shares: ShareTuple
    score: float
    evaluations: int


def _normalize_shares(s1: float, s2: float, min_share: float) -> ShareTuple | None:
    s3 = 1.0 - s1 - s2
    ms = max(min_share, 0.0)
    if s1 < ms - 1e-9 or s2 < ms - 1e-9 or s3 < ms - 1e-9:
        return None
    if abs(s1 + s2 + s3 - 1.0) > 1e-6:
        return None
    return (round(s1, 8), round(s2, 8), round(s3, 8))


def _start_points(min_share: float, grid_step: float) -> list[tuple[float, float]]:
    """Multi-start seeds: simplex center, corners, and a few grid_step hints."""
    ms = max(min_share, 0.0)
    max_pair = max(1.0 - 2 * ms, ms)
    center = round((1.0 - ms) / 3.0, 6)
    points: list[tuple[float, float]] = [
        (center, center),
        (ms, ms),
        (max_pair, ms),
        (ms, max_pair),
        (max_pair, max_pair),
    ]
    step = max(grid_step, 0.1)
    if ms <= step <= max_pair:
        points.append((step, ms))
        points.append((ms, step))
    deduped: list[tuple[float, float]] = []
    seen: set[tuple[float, float]] = set()
    for s1, s2 in points:
        if s1 + s2 > 1.0 - ms + 1e-9:
            continue
        key = (round(s1, 6), round(s2, 6))
        if key in seen:
            continue
        seen.add(key)
        deduped.append(key)
    return deduped


def optimize_shares_for_triplet(
    triplet: tuple[Component, Component, Component],
    score_fn: ScoreFn,
    *,
    min_share: float,
    grid_step: float = 0.1,
) -> OptimizeSharesResult | None:
    del triplet
    ms = max(min_share, 0.0)
    max_s1 = max(1.0 - 2 * ms, ms)
    evaluations = 0

    def objective(x: Any) -> float:
        nonlocal evaluations
        evaluations += 1
        shares = _normalize_shares(float(x[0]), float(x[1]), ms)
        if shares is None:
            return 1e18
        return score_fn(shares)

    bounds = [(ms, max_s1), (ms, max_s1)]
    constraints = [{"type": "ineq", "fun": lambda x: 1.0 - ms - float(x[0]) - float(x[1])}]

    best_score: float | None = None
    best_shares: ShareTuple | None = None

    for s1, s2 in _start_points(ms, grid_step):
        if s1 + s2 > 1.0 - ms + 1e-9:
            continue
        try:
            res = minimize(
                objective,
                [s1, s2],
                method="SLSQP",
                bounds=bounds,
                constraints=constraints,
                options={"ftol": 1e-8, "maxiter": 40},
            )
        except Exception:
            continue
        if not res.success and res.fun >= 1e17:
            continue
        shares = _normalize_shares(float(res.x[0]), float(res.x[1]), ms)
        if shares is None:
            continue
        score = float(res.fun)
        if best_score is None or score < best_score:
            best_score = score
            best_shares = shares

    if best_shares is None or best_score is None:
        return None
    return OptimizeSharesResult(shares=best_shares, score=best_score, evaluations=evaluations)