"""
Hybrid Dynamical System: Bouncing Ball in 1-D

This model demonstrates a hybrid system with:
- Flow State 1: Free fall (when ball center of mass is above radius r)
- Flow State 2: Spring-mass-damper (when ball is in contact with ground)
- Discrete transitions between these states
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp


class HybridBouncingBall:
    def __init__(self, m=0.1, r=0.1, g=9.81, k=5000.0, c=10.0):
        """
        Parameters:
        -----------
        m : float
            Mass of the ball (kg)
        r : float
            Radius of the ball (m)
        g : float
            Gravitational acceleration (m/s^2)
        k : float
            Spring constant when in contact with ground (N/m)
        c : float
            Damping coefficient when in contact with ground (N·s/m)
        """
        self.m = m
        self.r = r
        self.g = g
        self.k = k
        self.c = c

        # For tracking state transitions
        self.state_history = []
        self.transition_times = []

    def free_fall_dynamics(self, t, y):
        """
        Flow dynamics for free fall state.
        State: y = [position, velocity]
        """
        pos, vel = y
        dpos = vel
        dvel = -self.g
        return [dpos, dvel]

    def spring_damper_dynamics(self, t, y):
        """
        Flow dynamics for spring-mass-damper state.
        State: y = [position, velocity]

        When ball is compressed against ground:
        F = -k*(r - pos) - c*vel - m*g
        """
        pos, vel = y
        dpos = vel
        # Spring force kicks in when position < r
        # Compression is (r - pos)
        compression = self.r - pos
        spring_force = self.k * compression
        damping_force = self.c * vel

        dvel = (spring_force - damping_force - self.m * self.g) / self.m
        return [dpos, dvel]

    def event_contact_ground(self, t, y):
        """Event: Ball contacts ground (transition to spring-damper)"""
        pos, vel = y
        return pos - self.r

    def event_leave_ground(self, t, y):
        """Event: Ball leaves ground (transition to free fall)"""
        pos, vel = y
        # Leave ground when position > r AND velocity > 0
        if pos > self.r and vel > 0:
            return 0
        return 1

    # Make events terminal to stop integration
    event_contact_ground.terminal = True
    event_leave_ground.terminal = True

    def simulate(self, y0, t_span, max_transitions=20):
        """
        Simulate the hybrid system.

        Parameters:
        -----------
        y0 : list
            Initial state [position, velocity]
        t_span : tuple
            Time span (t_start, t_end)
        max_transitions : int
            Maximum number of state transitions to simulate

        Returns:
        --------
        t_all : array
            Time points
        y_all : array
            State trajectory
        states : list
            State labels ('free_fall' or 'spring_damper')
        """
        t_all = []
        y_all = []
        states = []

        current_state = "free_fall" if y0[0] > self.r else "spring_damper"
        current_y = y0
        current_t = t_span[0]
        t_end = t_span[1]

        transitions = 0

        while current_t < t_end and transitions < max_transitions:
            if current_state == "free_fall":
                # Integrate free fall until contact with ground
                sol = solve_ivp(
                    self.free_fall_dynamics,
                    [current_t, t_end],
                    current_y,
                    events=self.event_contact_ground,
                    dense_output=True,
                    max_step=0.01,
                )

                # Store results
                t_all.append(sol.t)
                y_all.append(sol.y.T)
                states.extend(["free_fall"] * len(sol.t))

                # Check if event occurred
                if sol.t_events[0].size > 0:
                    # Transition to spring-damper
                    current_state = "spring_damper"
                    current_t = sol.t[-1]
                    current_y = sol.y[:, -1]
                    self.transition_times.append(current_t)
                    transitions += 1
                else:
                    break

            else:  # spring_damper
                # Integrate spring-damper until leaving ground
                sol = solve_ivp(
                    self.spring_damper_dynamics,
                    [current_t, t_end],
                    current_y,
                    events=self.event_leave_ground,
                    dense_output=True,
                    max_step=0.01,
                )

                # Store results
                t_all.append(sol.t)
                y_all.append(sol.y.T)
                states.extend(["spring_damper"] * len(sol.t))

                # Check if event occurred
                if sol.t_events[0].size > 0:
                    # Transition to free fall
                    current_state = "free_fall"
                    current_t = sol.t[-1]
                    current_y = sol.y[:, -1]
                    self.transition_times.append(current_t)
                    transitions += 1
                else:
                    break

        # Concatenate all results
        t_all = np.concatenate(t_all)
        y_all = np.vstack(y_all)

        return t_all, y_all, states


def plot_simulation(t, y, states, ball, show_phase=True):
    """
    Plot the simulation results.

    Parameters:
    -----------
    t : array
        Time points
    y : array
        State trajectory
    states : list
        State labels
    ball : HybridBouncingBall
        Ball object
    show_phase : bool
        Whether to show phase portrait
    """
    # Convert states to numeric for coloring
    state_numeric = np.array([1 if s == "free_fall" else 2 for s in states])

    if show_phase:
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    else:
        fig, axes = plt.subplots(2, 1, figsize=(12, 8))
        axes = axes.reshape(-1, 1)

    # Plot 1: Position vs Time
    ax1 = axes[0, 0] if show_phase else axes[0]
    scatter = ax1.scatter(t, y[:, 0], c=state_numeric, s=1, cmap="coolwarm", alpha=0.6)
    ax1.axhline(
        y=ball.r,
        color="k",
        linestyle="--",
        label=f"Ground contact (h={ball.r}m)",
        linewidth=1,
    )
    ax1.axhline(
        y=0, color="gray", linestyle="-", label="Ground level", linewidth=1, alpha=0.5
    )

    # Mark transitions
    for t_trans in ball.transition_times:
        ax1.axvline(x=t_trans, color="green", linestyle=":", alpha=0.3, linewidth=1)

    ax1.set_xlabel("Time (s)", fontsize=11)
    ax1.set_ylabel("Position (m)", fontsize=11)
    ax1.set_title("Ball Position vs Time", fontsize=12, fontweight="bold")
    ax1.grid(True, alpha=0.3)
    ax1.legend(fontsize=9)

    # Add colorbar
    cbar1 = plt.colorbar(scatter, ax=ax1)
    cbar1.set_label("State", fontsize=10)
    cbar1.set_ticks([1.33, 1.67])
    cbar1.set_ticklabels(["Free Fall", "Spring-Damper"])

    # Plot 2: Velocity vs Time
    ax2 = axes[1, 0] if show_phase else axes[1]
    scatter2 = ax2.scatter(t, y[:, 1], c=state_numeric, s=1, cmap="coolwarm", alpha=0.6)

    # Mark transitions
    for t_trans in ball.transition_times:
        ax2.axvline(x=t_trans, color="green", linestyle=":", alpha=0.3, linewidth=1)

    ax2.axhline(y=0, color="k", linestyle="--", linewidth=1, alpha=0.3)
    ax2.set_xlabel("Time (s)", fontsize=11)
    ax2.set_ylabel("Velocity (m/s)", fontsize=11)
    ax2.set_title("Ball Velocity vs Time", fontsize=12, fontweight="bold")
    ax2.grid(True, alpha=0.3)

    # Add colorbar
    cbar2 = plt.colorbar(scatter2, ax=ax2)
    cbar2.set_label("State", fontsize=10)
    cbar2.set_ticks([1.33, 1.67])
    cbar2.set_ticklabels(["Free Fall", "Spring-Damper"])

    if show_phase:
        # Plot 3: Phase Portrait
        ax3 = axes[0, 1]
        scatter3 = ax3.scatter(
            y[:, 0], y[:, 1], c=state_numeric, s=1, cmap="coolwarm", alpha=0.6
        )
        ax3.axvline(
            x=ball.r, color="k", linestyle="--", label=f"Contact threshold", linewidth=1
        )
        ax3.axhline(y=0, color="gray", linestyle="-", linewidth=1, alpha=0.5)
        ax3.set_xlabel("Position (m)", fontsize=11)
        ax3.set_ylabel("Velocity (m/s)", fontsize=11)
        ax3.set_title("Phase Portrait", fontsize=12, fontweight="bold")
        ax3.grid(True, alpha=0.3)
        ax3.legend(fontsize=9)

        # Add colorbar
        cbar3 = plt.colorbar(scatter3, ax=ax3)
        cbar3.set_label("State", fontsize=10)
        cbar3.set_ticks([1.33, 1.67])
        cbar3.set_ticklabels(["Free Fall", "Spring-Damper"])

        # Plot 4: Energy
        ax4 = axes[1, 1]

        # Calculate energies
        KE = 0.5 * ball.m * y[:, 1] ** 2
        PE = ball.m * ball.g * y[:, 0]

        # Elastic potential energy when compressed
        elastic_PE = np.zeros_like(y[:, 0])
        for i, (pos, state) in enumerate(zip(y[:, 0], states)):
            if state == "spring_damper" and pos < ball.r:
                compression = ball.r - pos
                elastic_PE[i] = 0.5 * ball.k * compression**2

        total_E = KE + PE + elastic_PE

        ax4.plot(t, KE, label="Kinetic", linewidth=1.5, alpha=0.7)
        ax4.plot(t, PE, label="Gravitational PE", linewidth=1.5, alpha=0.7)
        ax4.plot(t, elastic_PE, label="Elastic PE", linewidth=1.5, alpha=0.7)
        ax4.plot(t, total_E, label="Total", linewidth=2, color="black", linestyle="--")

        # Mark transitions
        for t_trans in ball.transition_times:
            ax4.axvline(x=t_trans, color="green", linestyle=":", alpha=0.3, linewidth=1)

        ax4.set_xlabel("Time (s)", fontsize=11)
        ax4.set_ylabel("Energy (J)", fontsize=11)
        ax4.set_title("Energy Components", fontsize=12, fontweight="bold")
        ax4.grid(True, alpha=0.3)
        ax4.legend(fontsize=9)

    plt.tight_layout()
    return fig


if __name__ == "__main__":
    # Create ball with specific parameters
    ball = HybridBouncingBall(
        m=0.10,  # 1 kg mass
        r=0.1,  # 10 cm radius
        g=9.81,  # Earth gravity
        k=500.0,  # Spring constant
        c=0.89,  # Damping coefficient
    )

    # Initial conditions: drop from 2 meters with zero velocity
    y0 = [1.0, 0.0]  # [position (m), velocity (m/s)]

    # Simulate for 5 seconds
    t_span = (0, 5.0)

    print("Simulating hybrid bouncing ball system...")
    print(f"Initial conditions: h0 = {y0[0]} m, v0 = {y0[1]} m/s")
    print(
        f"Ball parameters: m={ball.m} kg, r={ball.r} m, k={ball.k} N/m, c={ball.c} N·s/m"
    )
    print()

    t, y, states = ball.simulate(y0, t_span, max_transitions=30)

    print(f"Simulation complete!")
    print(f"Total time simulated: {t[-1]:.3f} s")
    print(f"Number of state transitions: {len(ball.transition_times)}")
    print(f"Transition times: {[f'{tt:.3f}' for tt in ball.transition_times[:10]]}")
    print()

    # Count time in each state
    free_fall_count = states.count("free_fall")
    spring_damper_count = states.count("spring_damper")
    total_points = len(states)

    print(f"Time distribution:")
    print(f"  Free fall: {free_fall_count/total_points*100:.1f}%")
    print(f"  Spring-damper: {spring_damper_count/total_points*100:.1f}%")

    # Plot results
    fig = plot_simulation(t, y, states, ball, show_phase=True)
    plt.savefig(
        "/home/danesabo/Documents/Dane's Vault/Presentations/ERLM/bouncing_ball_hybrid.png",
        dpi=300,
        bbox_inches="tight",
    )
    print(f"\nPlot saved to: bouncing_ball_hybrid.png")
    plt.show()