import numpy as np
import torch


def _to_torch_dtype(dt):
    if isinstance(dt, torch.dtype) or dt is None:
        return dt or torch.get_default_dtype()

    mapping = {
        np.dtype("float32"): torch.float32,
        np.float32: torch.float32,
        np.dtype("float64"): torch.float64,
        np.float64: torch.float64,
        np.dtype("int32"): torch.int32,
        np.int32: torch.int32,
        np.dtype("int64"): torch.int64,
        np.int64: torch.int64,
    }
    return mapping.get(dt, torch.get_default_dtype())


def mb_params(device=None, dtype=torch.float32):

    dtype = _to_torch_dtype(dtype)
    A = torch.tensor([-200.0, -100.0, -170.0, 15.0], device=device, dtype=dtype)
    a = torch.tensor([-1.0, -1.0, -6.5, 0.7], device=device, dtype=dtype)
    b = torch.tensor([0.0, 0.0, 11.0, 0.6], device=device, dtype=dtype)
    c = torch.tensor([-10.0, -10.0, -6.5, 0.7], device=device, dtype=dtype)
    x0 = torch.tensor([1.0, 0.0, -0.5, -1.0], device=device, dtype=dtype)
    y0 = torch.tensor([0.0, 0.5, 1.5, 1.0], device=device, dtype=dtype)
    return A, a, b, c, x0, y0


def mb_potential_xy(x, y):
    """Compute Müller-Brown potential U(x,y). Args: x, y: torch tensors (broadcastable to same shape) Returns: U: torch tensor, same broadcasted shape"""
    device = (
        x.device
        if isinstance(x, torch.Tensor)
        else (y.device if isinstance(y, torch.Tensor) else None)
    )
    dt = (
        x.dtype
        if isinstance(x, torch.Tensor)
        else (y.dtype if isinstance(y, torch.Tensor) else None)
    )
    dt = _to_torch_dtype(dt)

    x = torch.as_tensor(x, dtype=dt, device=device)
    y = torch.as_tensor(y, dtype=dt, device=device)

    A, a, b, c, x0, y0 = mb_params(device=x.device, dtype=dt)
    X = x[..., None]
    Y = y[..., None]
    exp_arg = a * (X - x0) ** 2 + b * (X - x0) * (Y - y0) + c * (Y - y0) ** 2
    U = torch.sum(A * torch.exp(exp_arg), dim=-1)
    return U


def mb_grad_xy(x, y):
    """
    Gradient of U: returns (dU/dx, dU/dy)
    """
    device, dtype = x.device, x.dtype
    A, a, b, c, x0, y0 = mb_params(device=device, dtype=dtype)
    X = x[..., None]
    Y = y[..., None]
    dx = X - x0
    dy = Y - y0
    exp_arg = a * dx * dx + b * dx * dy + c * dy * dy
    E = torch.exp(exp_arg) * A
    dUdx = torch.sum(E * (2 * a * dx + b * dy), dim=-1)
    dUdy = torch.sum(E * (b * dx + 2 * c * dy), dim=-1)
    return dUdx, dUdy


def simulate_overdamped_mb(
    n_steps=100000,
    dt=1e-3,
    beta=1.0,
    mu=1.0,
    x0=None,
    y0=None,
    noise=True,
    device="cpu",
    dtype=torch.float32,
    seed=None,
):
    """
    Generate unconditioned trajectories on the MB surface.
    Overdamped Langevin: dx = -mu * grad U dt + sqrt(2 mu / beta) dW
    Returns numpy arrays (xs, ys) of shape [n_steps].
    A sequence of coordinates with a fixed dt
    """
    if seed is not None:
        torch.manual_seed(seed)
        np.random.seed(seed)
    x = torch.tensor([x0], device=device, dtype=dtype)
    y = torch.tensor([y0], device=device, dtype=dtype)
    xs = torch.empty(n_steps, dtype=dtype, device=device)
    ys = torch.empty(n_steps, dtype=dtype, device=device)
    sigma = (2.0 * mu / beta * dt) ** 0.5
    for t in range(n_steps):
        dUdx, dUdy = mb_grad_xy(x, y)
        x = x - mu * dUdx * dt + sigma * torch.randn_like(x)
        y = y - mu * dUdy * dt + sigma * torch.randn_like(y)
        xs[t] = x
        ys[t] = y
    return xs.cpu().numpy(), ys.cpu().numpy()