"""Landscape template library (v1 template-picker LandscapeForge).

Each template has a hand-written analytic gradient — no autodiff required.
All templates are guaranteed differentiable, finite, and bounded below on
the typical init region x ~ N(0, 0.5^2 I).
"""

from dataclasses import dataclass, field
from typing import Callable, Literal

import numpy as np


TemplateName = Literal[
    "quadratic", "styblinski_tang", "huber",
    "gaussian_mix", "himmelblau",
    "rosenbrock", "stiff_quadratic", "plateau", "cliff",
]

Tier = Literal["T0", "T1", "T2"]

TIER_MENU: dict[str, list[str]] = {
    "T0": ["quadratic", "styblinski_tang", "huber"],
    "T1": ["quadratic", "styblinski_tang", "huber", "gaussian_mix", "himmelblau"],
    "T2": ["quadratic", "styblinski_tang", "huber", "gaussian_mix", "himmelblau",
           "rosenbrock", "stiff_quadratic", "plateau", "cliff"],
}


@dataclass
class Landscape:
    name: str
    dim: int
    params: dict
    f: Callable[[np.ndarray], float]
    grad: Callable[[np.ndarray], np.ndarray]
    f_min: float = 0.0            # known global minimum value
    description: str = ""


# ---------- template constructors ----------

def make_quadratic(dim: int, cond: float = 1.0, rng: np.random.Generator | None = None) -> Landscape:
    """f(x) = 0.5 * x^T A x with diag Hessian of condition number `cond`."""
    diag = np.linspace(1.0, float(cond), dim)
    A = np.diag(diag)

    def f(x): return float(0.5 * x @ A @ x)
    def grad(x): return A @ x

    return Landscape(
        name="quadratic", dim=dim, params={"cond": cond},
        f=f, grad=grad, f_min=0.0,
        description=f"Convex quadratic in R^{dim}, condition number {cond:.1f}.",
    )


def make_stiff_quadratic(dim: int, cond: float = 1000.0, **_) -> Landscape:
    return make_quadratic(dim, cond)  # alias with higher cond


def make_styblinski_tang(dim: int, **_) -> Landscape:
    """f(x) = 0.5 * sum(x^4 - 16 x^2 + 5 x), min at x_i ≈ -2.903534."""
    def f(x):
        return float(0.5 * np.sum(x**4 - 16.0 * x**2 + 5.0 * x))

    def grad(x):
        return 0.5 * (4.0 * x**3 - 32.0 * x + 5.0)

    f_min = dim * 0.5 * ((-2.903534)**4 - 16.0 * (-2.903534)**2 + 5.0 * (-2.903534))
    return Landscape(
        name="styblinski_tang", dim=dim, params={},
        f=f, grad=grad, f_min=float(f_min),
        description=f"Styblinski-Tang in R^{dim}, multimodal with global min at x_i ≈ -2.9.",
    )


def make_huber(dim: int, delta: float = 1.0, **_) -> Landscape:
    """Smooth Huber-ish loss: f(x) = sum(delta^2 * (sqrt(1 + (x/delta)^2) - 1)).

    Smooth everywhere (unlike piecewise Huber). Behaves like 0.5 x^2 near 0,
    linear for |x| >> delta.
    """
    def f(x):
        return float(np.sum(delta**2 * (np.sqrt(1.0 + (x/delta)**2) - 1.0)))

    def grad(x):
        return x / np.sqrt(1.0 + (x/delta)**2)

    return Landscape(
        name="huber", dim=dim, params={"delta": delta},
        f=f, grad=grad, f_min=0.0,
        description=f"Smooth pseudo-Huber in R^{dim}, delta={delta}.",
    )


def make_rosenbrock(dim: int, **_) -> Landscape:
    """Classic stiff-valley Rosenbrock."""
    assert dim >= 2

    def f(x):
        return float(np.sum(100.0 * (x[1:] - x[:-1]**2)**2 + (1.0 - x[:-1])**2))

    def grad(x):
        g = np.zeros_like(x)
        g[:-1] += -400.0 * x[:-1] * (x[1:] - x[:-1]**2) - 2.0 * (1.0 - x[:-1])
        g[1:] += 200.0 * (x[1:] - x[:-1]**2)
        return g

    return Landscape(
        name="rosenbrock", dim=dim, params={},
        f=f, grad=grad, f_min=0.0,
        description=f"Rosenbrock (curved stiff valley) in R^{dim}, min at (1,..,1).",
    )


def make_gaussian_mix(dim: int, k: int = 3, sigma: float = 0.5, spread: float = 2.0,
                      rng: np.random.Generator | None = None, **_) -> Landscape:
    """f(x) = -sum_j w_j * exp(-||x - c_j||^2 / (2 sigma^2)).

    Negated so minima are where mixture density is highest. Bounded below by 0.
    """
    rng = rng if rng is not None else np.random.default_rng(0)
    centers = rng.normal(0.0, spread, size=(k, dim))
    weights = np.ones(k) / k   # uniform; one of these is the "global" min
    s2 = sigma * sigma

    def f(x):
        d2 = np.sum((centers - x)**2, axis=1)           # (k,)
        return float(-np.sum(weights * np.exp(-d2 / (2.0 * s2))))

    def grad(x):
        diffs = x - centers                              # (k, dim)
        d2 = np.sum(diffs**2, axis=1)                   # (k,)
        e = np.exp(-d2 / (2.0 * s2))                     # (k,)
        # d/dx [-w_j exp(-||x-c_j||^2 / 2s^2)] = w_j * (x-c_j) / s^2 * exp(...)
        return (weights * e / s2)[:, None] * diffs      # broadcast, sum over k below
        # Wait — need to sum over components:

    # Fix grad to properly sum:
    def grad_correct(x):
        diffs = x - centers
        d2 = np.sum(diffs**2, axis=1)
        e = np.exp(-d2 / (2.0 * s2))
        coeff = weights * e / s2
        return np.sum(coeff[:, None] * diffs, axis=0)

    # Global min approx: evaluate at each center, return the lowest f.
    f_min = float(min(f(c) for c in centers))
    return Landscape(
        name="gaussian_mix", dim=dim,
        params={"k": k, "sigma": sigma, "spread": spread, "centers": centers.tolist()},
        f=f, grad=grad_correct, f_min=f_min,
        description=f"Negative Gaussian mixture in R^{dim}, k={k} modes, sigma={sigma}, spread={spread}.",
    )


def make_himmelblau(dim: int = 2, **_) -> Landscape:
    """f(x,y) = (x^2+y-11)^2 + (x+y^2-7)^2. 4 global minima at value 0."""
    assert dim == 2, "Himmelblau is 2D only"

    def f(x):
        return float((x[0]**2 + x[1] - 11.0)**2 + (x[0] + x[1]**2 - 7.0)**2)

    def grad(x):
        gx = 4.0 * x[0] * (x[0]**2 + x[1] - 11.0) + 2.0 * (x[0] + x[1]**2 - 7.0)
        gy = 2.0 * (x[0]**2 + x[1] - 11.0) + 4.0 * x[1] * (x[0] + x[1]**2 - 7.0)
        return np.array([gx, gy])

    return Landscape(
        name="himmelblau", dim=2, params={}, f=f, grad=grad, f_min=0.0,
        description="Himmelblau (2D), four global minima at value 0.",
    )


def make_plateau(dim: int, radius: float = 1.0, **_) -> Landscape:
    """Smooth plateau: f = tanh((||x||^2 - r^2) / r^2). Near-zero gradient far from boundary."""
    r2 = radius * radius

    def f(x):
        return float(np.tanh((np.sum(x**2) - r2) / r2))

    def grad(x):
        u = (np.sum(x**2) - r2) / r2
        return (1.0 - np.tanh(u)**2) * (2.0 * x / r2)

    return Landscape(
        name="plateau", dim=dim, params={"radius": radius},
        f=f, grad=grad, f_min=-1.0,
        description=f"Plateau landscape in R^{dim}, radius {radius}, vanishing gradient at center.",
    )


def make_cliff(dim: int, **_) -> Landscape:
    """Smooth cliff: quadratic + tanh step. Tough for fixed-step optimizers."""
    def f(x):
        s = np.sum(x)
        return float(0.5 * np.sum(x**2) + 5.0 * np.tanh(s))

    def grad(x):
        s = np.sum(x)
        t = 1.0 - np.tanh(s)**2
        return x + 5.0 * t * np.ones_like(x)

    return Landscape(
        name="cliff", dim=dim, params={},
        f=f, grad=grad, f_min=-5.0,  # approximate lower bound
        description=f"Quadratic with tanh cliff in R^{dim}.",
    )


BUILDERS: dict[str, Callable[..., Landscape]] = {
    "quadratic": make_quadratic,
    "stiff_quadratic": make_stiff_quadratic,
    "styblinski_tang": make_styblinski_tang,
    "huber": make_huber,
    "rosenbrock": make_rosenbrock,
    "gaussian_mix": make_gaussian_mix,
    "himmelblau": make_himmelblau,
    "plateau": make_plateau,
    "cliff": make_cliff,
}


def build_landscape(template: str, dim: int, params: dict | None = None,
                    rng: np.random.Generator | None = None) -> Landscape:
    """Instantiate a landscape by name."""
    if template not in BUILDERS:
        raise ValueError(f"Unknown template {template!r}; known: {list(BUILDERS)}")
    return BUILDERS[template](dim=dim, rng=rng, **(params or {}))


def structural_hints(ls: Landscape, n_samples: int = 200,
                     rng: np.random.Generator | None = None) -> dict:
    """Env-computed hints: Lipschitz estimate, gradient spread, modality hint.

    Sampled at reset, exposed to OptCoder as free info.
    """
    rng = rng if rng is not None else np.random.default_rng(0)
    xs = rng.normal(0.0, 1.0, size=(n_samples, ls.dim))
    fs = np.array([ls.f(x) for x in xs])
    gs = np.array([ls.grad(x) for x in xs])
    g_norms = np.linalg.norm(gs, axis=1)
    return {
        "lipschitz_estimate": float(np.percentile(g_norms, 95)),
        "grad_norm_median": float(np.median(g_norms)),
        "f_range": [float(fs.min()), float(fs.max())],
        "f_median": float(np.median(fs)),
        # crude modality: count local f peaks on random 1D slices
        "modality_hint": _modality_hint(ls, rng),
    }


def _modality_hint(ls: Landscape, rng: np.random.Generator) -> str:
    """Very crude multimodality probe on 5 random 1D slices."""
    hits = 0
    for _ in range(5):
        center = rng.normal(0.0, 0.5, size=ls.dim)
        direction = rng.normal(0.0, 1.0, size=ls.dim)
        direction /= np.linalg.norm(direction) + 1e-12
        ts = np.linspace(-3.0, 3.0, 30)
        vals = np.array([ls.f(center + t * direction) for t in ts])
        # count sign changes in finite diff
        d = np.diff(vals)
        s = np.sign(d)
        sign_changes = int(np.sum(s[1:] != s[:-1]))
        if sign_changes >= 3:
            hits += 1
    if hits >= 3:
        return "multimodal"
    if hits >= 1:
        return "possibly_multimodal"
    return "unimodal"