"""Fuzzy controller and closed-loop cart-pendulum simulation from the article.

Requires NumPy. Run with: python fuzzy_control.py
"""

import numpy as np

def gaussian(x, mu=0., sig=1.):
    return np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))


def fuzzy_or(a, b, type="max"):
    if type == "max":
        return np.fmax(a, b)


def fuzzy_or_array(a, type="max"):
    if type == "max":
        return np.max(a, axis=0)


def fuzzy_and(a, b, type="min"):
    if type == "min":
        return np.fmin(a, b)



class FuzzyControl:
    def __init__(self):
        self.resolution = 50
        self.r = 30
        self.sample_points = np.linspace(-self.r, self.r, self.resolution)

    def fuzzify(self, angle, angular_speed):
        angle_memberships = []
        angular_speed_memberships = []
        #
        if angle > 50:
            angle = 50
        if angle < -50:
            angle = -50
        if angular_speed > 45:
            angular_speed = 45
        if angular_speed < -45:
            angular_speed = -45

        angle_mu = [-50., -18., -3., 3., 18., 50.]
        angle_sig = [15., 9.5, 2.5, 2.5, 9.5, 15.]
        for i in range(6):
            angle_memberships.append(gaussian(angle, angle_mu[i], sig=angle_sig[i]))

        angular_speed_mu = [-45., -15., -3., 3., 15., 45.]
        angular_speed_sig = [14., 7.5, 3., 3., 7.5, 14.]
        for i in range(6):
            angular_speed_memberships.append(gaussian(angular_speed, angular_speed_mu[i], sig=angular_speed_sig[i]))
        return angle_memberships, angular_speed_memberships

    def inference(self, angle_memberships, angular_speed_membership):
        #
        # How to interpret the rule base:
        # NB: negative big, NM: negative medium, NS: negative small
        # PB: positive big, PM: positive medium, PS: positive small
        # The table:
        #    NB, NM, NS, PS, PM, PB angular speed
        # NB  0   0   0   1   2   3
        # NM  0   0   1   2   3   3
        # NS  0   1   2   3   3   4
        # PS  1   2   2   3   4   5
        # PM  2   2   3   4   5   5
        # PB  2   3   4   5   5   5
        # angle
        #
        # One cell contains the control force:
        # 0 NB, 1 NM, 2 NS, 3 PS, 4 PM, 5 PB
        rule_base = [[0, 0, 0, 1, 2, 3],
                     [0, 0, 1, 2, 3, 3],
                     [0, 1, 2, 3, 3, 4],
                     [1, 2, 2, 3, 4, 5],
                     [2, 2, 3, 4, 5, 5],
                     [2, 3, 4, 5, 5, 5]]

        fuzzy_control_force = [0., 0., 0., 0., 0., 0.]
        for i in range(6):
            for k in range(6):
                # AND joins the two conditions; OR combines rules for a force term.
                pt = fuzzy_and(angle_memberships[i], angular_speed_membership[k])
                fuzzy_control_force[rule_base[i][k]] = fuzzy_or(fuzzy_control_force[rule_base[i][k]], pt)

        for i, mu in enumerate([h * 10 for h in [-3, -2., -1., 1., 2., 3]]):
            fuzzy_control_force[i] = fuzzy_and(gaussian(self.sample_points, mu, sig=2.), fuzzy_control_force[i])
        fuzzy_control_force = fuzzy_or_array(np.asarray(fuzzy_control_force))
        return fuzzy_control_force

    def defuzzify(self, fuzzy_control_force):
        rv = (np.sum(self.sample_points * fuzzy_control_force) + 1e-8) / (np.sum(fuzzy_control_force) + 1e-8)
        return rv


# Cart and uniform rod from the previous post, in SI units.
# phi is the tilt from upright in radians, positive when the rod leans right.
M, m, L = 1.0, 0.5, 1.0       # cart mass (kg), rod mass (kg), rod length (m)
b, g = 0.1, 9.81              # cart damping (N s/m), gravity (m/s^2)
l = L / 2                     # pivot to the rod's center of mass (m)
J = m * L**2 / 12 + m * l**2  # rod's moment of inertia about the pivot (kg m^2)


def step(state, force, dt):
    """Advance (x, v, phi, omega) by one velocity-first Euler step."""
    x, v, phi, omega = state
    a = M + m
    coupling = m * l * np.cos(phi)
    rhs_x = force - b * v + m * l * omega**2 * np.sin(phi)
    rhs_phi = m * g * l * np.sin(phi)
    determinant = a * J - coupling**2
    cart_acc = (J * rhs_x - coupling * rhs_phi) / determinant
    angular_acc = (a * rhs_phi - coupling * rhs_x) / determinant
    v += cart_acc * dt
    omega += angular_acc * dt
    return x + v * dt, v, phi + omega * dt, omega


def control_force(controller, phi, omega):
    """Convert the state to degrees and run the three fuzzy stages."""
    angle = (np.degrees(phi) + 180.) % 360. - 180.
    angle_memberships, speed_memberships = controller.fuzzify(angle, np.degrees(omega))
    force_memberships = controller.inference(angle_memberships, speed_memberships)
    return controller.defuzzify(force_memberships)


def simulate(controller, state, duration, dt=0.001):
    """Apply the fuzzy force at every step and return the visited states."""
    states = [state]
    for _ in range(round(duration / dt)):
        force = control_force(controller, state[2], state[3])
        state = step(state, force, dt)
        states.append(state)
    return states


if __name__ == "__main__":
    controller = FuzzyControl()
    angle_memberships, speed_memberships = controller.fuzzify(-5., -3.)
    force_memberships = controller.inference(angle_memberships, speed_memberships)
    force = controller.defuzzify(force_memberships)
    print(f"Control force: {force:.3f} N")

    controller = FuzzyControl()
    states = simulate(controller, (0., 0., np.radians(-5.), np.radians(-3.)), duration=10.)
    x, v, phi, omega = states[-1]
    print(f"After 10 s: tilt {np.degrees(phi):.2f} deg, cart at {x:.1f} m moving {v:.2f} m/s")