"""Desktop animation of the fuzzy-controlled cart and rod.

Requires NumPy and Arcade 3, with fuzzy_control.py in the same directory:
    python -m pip install numpy arcade
    python arcade_simulation.py
Controls: hold left/right arrows to push with 10 N, Enter toggles the controller,
R resets, and Space saves screenshot.png.
"""

import math

import arcade
import numpy as np

from fuzzy_control import FuzzyControl, control_force, step

SCREEN_WIDTH = 1200
SCREEN_HEIGHT = 600
PIXELS_PER_METER = 200
PUSH = 10.  # N
DT = 0.001  # s, physics timestep
INITIAL_STATE = (0., 0., math.radians(-5.), math.radians(-3.))  # x, v, phi, omega


class Application(arcade.Window):
    def __init__(self):
        super().__init__(SCREEN_WIDTH, SCREEN_HEIGHT, title="Fuzzy control of a cart and rod")
        self.background_color = arcade.color.WHITE
        self.controller = FuzzyControl()
        self.texts = [arcade.Text("", SCREEN_WIDTH - 10, 10 + 24 * i, arcade.color.BLACK, 14, anchor_x="right")
                      for i in range(6)]
        self.curve_label = arcade.Text("Force memberships before defuzzification", 10, 10,
                                       arcade.color.BLACK, 12)
        self.reset()

    def reset(self):
        self.state = INITIAL_STATE
        self.sim_time = 0.
        self.accumulator = 0.
        self.push = 0.
        self.controller_on = True

    def force_memberships(self):
        angle = (math.degrees(self.state[2]) + 180.) % 360. - 180.
        memberships = self.controller.fuzzify(angle, math.degrees(self.state[3]))
        return self.controller.inference(*memberships)

    def on_update(self, delta_time):
        # Advance the physics in fixed steps so the motion follows real time.
        self.accumulator += min(delta_time, 0.05)
        while self.accumulator >= DT:
            force = self.push
            if self.controller_on:
                force += control_force(self.controller, self.state[2], self.state[3])
            self.state = step(self.state, force, DT)
            self.sim_time += DT
            self.accumulator -= DT

    def draw_arrow(self, x, y, force, color):
        length = 3 * force  # 3 px per newton
        if abs(length) < 1:
            return
        head = math.copysign(min(10, abs(length)), length)
        arcade.draw_line(x, y, x + length, y, color, 4)
        arcade.draw_line(x + length, y, x + length - head, y + 7, color, 4)
        arcade.draw_line(x + length, y, x + length - head, y - 7, color, 4)

    def on_draw(self):
        self.clear()
        x, v, phi, omega = self.state
        # The camera follows the cart; the track marks show how far it has moved.
        pivot_x, pivot_y = SCREEN_WIDTH / 2, 300
        track_y = pivot_y - 45
        arcade.draw_line(0, track_y, SCREEN_WIDTH, track_y, arcade.color.BLACK, 2)
        half_view = SCREEN_WIDTH / (2 * PIXELS_PER_METER)
        for mark in range(math.ceil(x - half_view), math.floor(x + half_view) + 1):
            mark_x = pivot_x + (mark - x) * PIXELS_PER_METER
            arcade.draw_line(mark_x, track_y, mark_x, track_y - 10, arcade.color.GRAY, 2)

        arcade.draw_lrbt_rectangle_outline(pivot_x - 100, pivot_x + 100, pivot_y - 25, pivot_y + 25,
                                           arcade.color.BLACK, 2)
        for dx in (-60, 60):
            arcade.draw_circle_filled(pivot_x + dx, pivot_y - 25, 20, arcade.color.BLACK)
        tip_x = pivot_x + PIXELS_PER_METER * math.sin(phi)
        tip_y = pivot_y + PIXELS_PER_METER * math.cos(phi)
        arcade.draw_line(pivot_x, pivot_y, tip_x, tip_y, arcade.color.RED, 5)
        arcade.draw_circle_filled(pivot_x, pivot_y, 10, arcade.color.BLACK)

        memberships = self.force_memberships()
        force = self.controller.defuzzify(memberships)
        applied = force if self.controller_on else 0.
        self.draw_arrow(pivot_x, pivot_y + 8, applied, arcade.color.BLUE)
        self.draw_arrow(pivot_x, pivot_y - 8, self.push, arcade.color.ORANGE)

        # Combined force memberships over -30 to 30 N, drawn 400 px wide and 150 px tall.
        points = [(10 + 400 * (u + 30) / 60, 35 + 150 * mu)
                  for u, mu in zip(self.controller.sample_points, memberships)]
        arcade.draw_line_strip(points, arcade.color.AMBER, 2)
        centroid_x = 10 + 400 * (force + 30) / 60
        arcade.draw_line(centroid_x, 35, centroid_x, 185, arcade.color.BLUE, 2)
        self.curve_label.draw()

        status = "on" if self.controller_on else "off"
        lines = [f"{np.degrees(omega):.2f}°/s angular velocity",
                 f"{np.degrees(phi):.2f}° angle",
                 f"{v:.2f} m/s cart velocity",
                 f"{self.push:.0f} N push",
                 f"{force:.2f} N controller force ({status})",
                 f"{self.sim_time:.1f} s"]
        for text, value in zip(self.texts, lines):
            text.text = value
            text.draw()

    def on_key_press(self, key, modifiers):
        if key == arcade.key.LEFT:
            self.push = -PUSH
        elif key == arcade.key.RIGHT:
            self.push = PUSH
        elif key == arcade.key.ENTER:
            self.controller_on = not self.controller_on
        elif key == arcade.key.R:
            self.reset()
        elif key == arcade.key.SPACE:
            arcade.get_image().save("screenshot.png", "PNG")

    def on_key_release(self, key, modifiers):
        if key in (arcade.key.LEFT, arcade.key.RIGHT):
            self.push = 0.


if __name__ == "__main__":
    Application()
    arcade.run()