from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, cm


# -----------------------------
# Configuration
# -----------------------------

@dataclass(frozen=True)
class SimConfig:
    simulation_time_s: float = 4.0
    sim_dt_s: float = 0.001
    viz_dt_s: float = 0.1  # animation timestep (seconds)

    # DP parameters
    dp_resolution: int = 51
    control_dt_s: float = 0.1
    num_control_options: int = 10
    control_penalty: float = 0.05

    # Convergence
    max_dp_iters: int = 10_000
    policy_change_tol: float = 0.01  # <=1% cells change => converged


@dataclass(frozen=True)
class PlantParams:
    m1: float = 1.0
    l1: float = 1.0
    g: float = 9.81


# -----------------------------
# Plant + integration (behavior-preserving)
# -----------------------------

def plant_ddtheta(theta: float, u: float, params: PlantParams) -> float:
    """Angular acceleration (ddtheta) for the single pendulum."""
    return (params.g / params.l1) * math.sin(theta) + u


def simulation_advance(state: np.ndarray, ddtheta: float, dt: float) -> np.ndarray:
    """
    Behavior-preserving integrator (matches your original code).

    state = [theta, omega, alpha]
    but uses omega = state[2] (previous alpha) exactly like your original:
        omega1 = states[2]
        alpha1 = ddtheta
        theta += omega1 * dt
        omega1 += alpha1 * dt
        next = [theta, omega1, alpha1]
    """
    theta = float(state[0])
    omega = float(state[2])     # NOTE: intentional to preserve existing behavior
    alpha = float(ddtheta)

    theta = theta + omega * dt
    omega = omega + alpha * dt

    return np.array([theta, omega, alpha], dtype=float)


# -----------------------------
# Dynamic Programming Controller
# -----------------------------

@dataclass
class DPConfig:
    theta_lim: Tuple[float, float] = (-2 * math.pi, 2 * math.pi)
    omega_lim: Tuple[float, float] = (-10.0, 10.0)


class DynamicProgramming:
    def __init__(self, cfg: DPConfig = DPConfig()):
        self.cfg = cfg

        # Filled during get_control_policy(...)
        self.plant: PlantParams | None = None
        self.set_point: np.ndarray | None = None
        self.resolution: int | None = None
        self.control_dt: float | None = None
        self.control_options: np.ndarray | None = None

        self.policy_map: np.ndarray | None = None
        self.state_cost: np.ndarray | None = None

    # ---- public API ----

    def get_control_policy(
        self,
        plant: PlantParams,
        initial_state: np.ndarray,
        set_point: np.ndarray,
        resolution: int,
        control_dt: float,
        num_control_options: int,
        control_penalty: float,
        plotting: bool = False,
        max_iters: int = 10_000,
        policy_change_tol: float = 0.01,
    ) -> np.ndarray:
        """
        Compute policy_map(theta_id, omega_id) -> u via DP iterations.
        Preserves your original update rule:
          cost_candidate = last_cost[s] + last_cost[s'] + penalty*|u|
        and convergence: <1% of policy entries changed.
        """
        self.plant = plant
        self.set_point = np.array(set_point, dtype=float)
        self.resolution = int(resolution)
        self.control_dt = float(control_dt)
        self.control_options = np.linspace(-50.0, 50.0, num=num_control_options)

        self.policy_map = np.zeros((self.resolution, self.resolution), dtype=float)
        self.state_cost = self._init_state_cost()
        self._zero_set_point_cost()

        if plotting:
            self._plot_surface(self.state_cost, title="Initial state cost surface", zlabel="Cost")

        converged = False
        last_cost = self.state_cost.copy()
        last_policy = self.policy_map.copy()

        candidates = np.empty(self.control_options.shape[0], dtype=float)

        for _ in range(max_iters):
            # full sweep update
            for t_id in range(self.resolution):
                for o_id in range(self.resolution):
                    # evaluate each candidate control
                    for u_idx, u in enumerate(self.control_options):
                        t_next, o_next = self._advance_dynamics(t_id, o_id, float(u))
                        candidates[u_idx] = (
                            last_cost[t_id, o_id]
                            + last_cost[t_next, o_next]
                            + control_penalty * abs(float(u))
                        )

                    best = int(np.argmin(candidates))
                    self.policy_map[t_id, o_id] = float(self.control_options[best])
                    self.state_cost[t_id, o_id] = float(candidates[best])

            # convergence check (same criterion, but clearer)
            changed = (self.policy_map != last_policy)
            frac_changed = float(np.mean(changed))
            if frac_changed <= policy_change_tol:
                converged = True
                break

            last_policy = self.policy_map.copy()
            last_cost = self.state_cost.copy()

        if not converged:
            print("Warning: DP did not converge within max_iters; using last policy.")

        if plotting:
            self._plot_surface(self.policy_map, title="Learned policy surface", zlabel="u")

        return self.policy_map

    def action(self, theta: float, omega: float) -> float:
        """Lookup control u from the discretized policy."""
        if self.policy_map is None:
            raise RuntimeError("Policy not computed yet.")
        t_id = self.state_value_to_id(theta, "t1")
        o_id = self.state_value_to_id(omega, "o1")
        return float(self.policy_map[t_id, o_id])

    # ---- discretization helpers ----

    def state_id_to_value(self, idx: int, which: str) -> float:
        if self.resolution is None:
            raise RuntimeError("DP not initialized.")
        if which == "t1":
            lo, hi = self.cfg.theta_lim
        elif which == "o1":
            lo, hi = self.cfg.omega_lim
        else:
            raise ValueError("which must be 't1' or 'o1'")

        return idx * (hi - lo) / self.resolution + lo

    def state_value_to_id(self, val: float, which: str) -> int:
        if self.resolution is None:
            raise RuntimeError("DP not initialized.")
        if which == "t1":
            lo, hi = self.cfg.theta_lim
        elif which == "o1":
            lo, hi = self.cfg.omega_lim
        else:
            raise ValueError("which must be 't1' or 'o1'")

        if val < lo:
            return 0
        if val > hi:
            return self.resolution - 1

        # match original formula
        return int((val - lo) * self.resolution / (hi - lo))

    # ---- internal methods ----

    def _advance_dynamics(self, t_id: int, o_id: int, u: float) -> Tuple[int, int]:
        """Advance one control step in discretized space (same as your advance_dynamics)."""
        assert self.plant is not None
        assert self.control_dt is not None

        theta = self.state_id_to_value(t_id, "t1")
        omega = self.state_id_to_value(o_id, "o1")

        theta_next = theta + omega * self.control_dt
        alpha = (-(self.plant.g / self.plant.l1) * math.sin(theta) + u)
        omega_next = omega + alpha * self.control_dt

        t_next = self.state_value_to_id(theta_next, "t1")
        o_next = self.state_value_to_id(omega_next, "o1")
        return t_next, o_next

    def _init_state_cost(self) -> np.ndarray:
        """Initialize cost surface as (theta-theta_d)^2 + (omega-omega_d)^2."""
        assert self.resolution is not None
        assert self.set_point is not None

        cost = np.zeros((self.resolution, self.resolution), dtype=float)
        theta_d, omega_d = float(self.set_point[0]), float(self.set_point[1])

        for t_id in range(self.resolution):
            theta = self.state_id_to_value(t_id, "t1")
            ct = (theta - theta_d) ** 2
            for o_id in range(self.resolution):
                omega = self.state_id_to_value(o_id, "o1")
                co = (omega - omega_d) ** 2
                cost[t_id, o_id] = ct + co

        return cost

    def _zero_set_point_cost(self) -> None:
        """Force exact setpoint cell to zero cost (same intent as your zero_set_point)."""
        assert self.resolution is not None
        assert self.set_point is not None

        theta_d, omega_d = float(self.set_point[0]), float(self.set_point[1])
        t_id = int((theta_d - self.cfg.theta_lim[0]) / (self.cfg.theta_lim[1] - self.cfg.theta_lim[0]) * self.resolution)
        o_id = int((omega_d - self.cfg.omega_lim[0]) / (self.cfg.omega_lim[1] - self.cfg.omega_lim[0]) * self.resolution)

        # bounds safety
        t_id = max(0, min(self.resolution - 1, t_id))
        o_id = max(0, min(self.resolution - 1, o_id))

        self.state_cost[t_id, o_id] = 0.0

    @staticmethod
    def _plot_surface(Z: np.ndarray, title: str, zlabel: str) -> None:
        """3D surface plot helper (same plots as before, cleaner)."""
        fig = plt.figure()
        ax = fig.add_subplot(projection="3d")
        # Mesh just for visualization; values are Z[y,x]-like but we keep your convention
        X = np.arange(Z.shape[0])
        Y = np.arange(Z.shape[1])
        X, Y = np.meshgrid(X, Y)
        surf = ax.plot_surface(X, Y, Z.T, cmap=cm.coolwarm, linewidth=0, antialiased=False)
        ax.set_title(title)
        ax.set_xlabel("theta_id")
        ax.set_ylabel("omega_id")
        ax.set_zlabel(zlabel)
        fig.colorbar(surf, shrink=0.5, aspect=5)
        plt.show()


# -----------------------------
# Animation
# -----------------------------

def animation_update(frame_idx: int, state_history: np.ndarray, l1: float, line):
    theta = float(state_history[0, frame_idx])
    x = np.array([0.0, l1 * math.sin(theta)])
    y = np.array([0.0, -l1 * math.cos(theta)])
    line.set_data(x, y)
    return (line,)


# -----------------------------
# Main
# -----------------------------

def run():
    cfg = SimConfig()
    plant = PlantParams()

    sim_steps = int(cfg.simulation_time_s / cfg.sim_dt_s)
    frame_stride = max(1, int(cfg.viz_dt_s / cfg.sim_dt_s))

    # state = [theta, omega, alpha]
    state_history = np.zeros((3, sim_steps), dtype=float)
    state_history[:, 0] = np.array([0.0, 0.0, 0.0], dtype=float)

    set_point = np.array([math.pi, 0.0], dtype=float)

    controller = DynamicProgramming()
    controller.get_control_policy(
        plant=plant,
        initial_state=state_history[:, 0],
        set_point=set_point,
        resolution=cfg.dp_resolution,
        control_dt=cfg.control_dt_s,
        num_control_options=cfg.num_control_options,
        control_penalty=cfg.control_penalty,
        plotting=True,
        max_iters=cfg.max_dp_iters,
        policy_change_tol=cfg.policy_change_tol,
    )

    # simulate with policy lookup
    for i in range(1, sim_steps):
        theta_prev = float(state_history[0, i - 1])
        omega_prev = float(state_history[1, i - 1])
        u = controller.action(theta_prev, omega_prev)

        ddtheta = plant_ddtheta(theta_prev, u, plant)
        state_history[:, i] = simulation_advance(state_history[:, i - 1], ddtheta, cfg.sim_dt_s)

    # animate + export
    fig = plt.figure()
    ax = plt.axes(xlim=(-plant.l1, plant.l1), ylim=(-plant.l1, plant.l1))
    line, = ax.plot([], [], "k-o")

    frames = range(0, sim_steps, frame_stride)
    ani = animation.FuncAnimation(
        fig,
        animation_update,
        frames=frames,
        fargs=(state_history, plant.l1, line),
        blit=True,
        save_count=int(cfg.simulation_time_s / cfg.viz_dt_s),
        repeat=False,
    )
    plt.show()

    html = ani.to_jshtml(fps=1 / cfg.viz_dt_s)
    with open("SP_dynamicprogramming.html", "w", encoding="utf-8") as f:
        f.write(html)


if __name__ == "__main__":
    run()