import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import gradio as gr

def parse_phi(phi_str):
    try:
        return float(eval(phi_str.replace("pi","np.pi")))
    except:
        return 0.0

def add_watermark(ax, text="AerovFX"):
    ax.text(0.95, 0.02, text, transform=ax.transAxes,
            fontsize=10, color='gray', alpha=0.5,
            ha='right', va='bottom')

def setup_ax(A, t_max):
    fig, ax = plt.subplots(figsize=(6,3))
    ax.set_xlim(0, t_max)
    ax.set_ylim(-1.2*A, 1.2*A)
    
    # Ox: mỗi ô 0.1 s
    ax.set_xticks(np.arange(0, t_max+0.1, 0.1))
    # Oy: mỗi ô A/10
    y_step = A / 10 if A>0 else 0.2
    ax.set_yticks(np.arange(-1.2*A, 1.2*A+y_step, y_step))
    
    # Grid
    ax.grid(which='major', color='blue', linestyle='-', alpha=0.5)
    ax.grid(which='minor', color='blue', linestyle='-', alpha=0.3)
    
    # Trục tọa độ
    ax.spines['left'].set_position('zero')
    ax.spines['bottom'].set_position('zero')
    ax.spines['left'].set_color('black'); ax.spines['bottom'].set_color('black')
    ax.spines['left'].set_linewidth(2); ax.spines['bottom'].set_linewidth(2)
    ax.spines['right'].set_color('none'); ax.spines['top'].set_color('none')
    
    # Mũi tên Ox góc phải
    ax.annotate("", xy=(t_max,0), xytext=(0,0), arrowprops=dict(arrowstyle="->", lw=2, color='black'))
    ax.annotate("", xy=(0,1.2*A), xytext=(0,0), arrowprops=dict(arrowstyle="->", lw=2, color='black'))
    ax.text(t_max, -0.1*A, "t (s)", ha='right', va='top', fontsize=6)
    ax.text(0.05, 1.05*A, "x (cm)", ha='left', va='bottom', fontsize=6)
    ax.annotate("A", xy=(0,A), xytext=(-0.05,A), textcoords='data', color='red', fontsize=6)
    
    # Giảm cỡ chữ số trên trục
    ax.tick_params(axis='both', which='major', labelsize=6)
    ax.tick_params(axis='both', which='minor', labelsize=5)
    
    add_watermark(ax)
    return fig, ax

def generate_gif(A, f, omega, phi_str, t_max):
    phi = parse_phi(phi_str)
    fps = 10
    if f and f>0:
        omega = 2*np.pi*f
    elif omega and omega>0:
        f = omega/(2*np.pi)
    else:
        return "⚠️ Cần nhập f>0 hoặc ω>0", None

    n_frames = int(t_max * fps)
    t = np.linspace(0, t_max, 500)
    fig, ax = setup_ax(A, t_max)
    line, = ax.plot([], [], 'b-', lw=2)

    def init(): line.set_data([], []); return line,
    def update(frame):
        t_current = frame / fps
        y = A * np.cos(omega * t + phi + omega * t_current)
        line.set_data(t, y)
        return line,

    ani = animation.FuncAnimation(fig, update, frames=n_frames, init_func=init, blit=True)
    gif_path = "cosine.gif"
    ani.save(gif_path, writer='pillow', fps=fps)
    plt.close(fig)
    return f"f={f:.3f}Hz, ω={omega:.3f}rad/s, φ={phi:.3f}rad", gif_path

def generate_mp4(A, f, omega, phi_str, t_max):
    phi = parse_phi(phi_str)
    fps = 30
    if f and f>0:
        omega = 2*np.pi*f
    elif omega and omega>0:
        f = omega/(2*np.pi)
    else:
        return "⚠️ Cần nhập f>0 hoặc ω>0", None

    n_frames = int(t_max * fps)
    t = np.linspace(0, t_max, 500)
    fig, ax = setup_ax(A, t_max)
    line, = ax.plot([], [], 'b-', lw=2)

    def init(): line.set_data([], []); return line,
    def update(frame):
        t_current = frame / fps
        y = A * np.cos(omega * t + phi + omega * t_current)
        line.set_data(t, y)
        return line,

    ani = animation.FuncAnimation(fig, update, frames=n_frames, init_func=init, blit=True)
    mp4_path = "cosine.mp4"
    ani.save(mp4_path, writer='ffmpeg', fps=fps, dpi=150)
    plt.close(fig)
    return f"f={f:.3f}Hz, ω={omega:.3f}rad/s, φ={phi:.3f}rad", mp4_path

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 📈 Đồ thị dao động điều hoà.")
    with gr.Row():
        with gr.Column(scale=1):
            A = gr.Number(value=2.0, label="Biên độ A (cm)")
            f = gr.Number(value=None, label="Tần số f (Hz)")
            omega = gr.Number(value=None, label="Tần số góc ω (rad/s)", info="Nhập tần số f để ô này tự ẩn")
            phi = gr.Textbox(value="0", label="Pha ban đầu φ (rad, ví dụ: 0, pi/2, pi)")
            t_max = gr.Number(value=2, label="Thời gian mô phỏng (s)")

        with gr.Column(scale=2):
            btn_gif = gr.Button("Tạo GIF")
            out_gif = gr.Image(type="filepath", label="GIF dao động")

        with gr.Column(scale=2):
            btn_mp4 = gr.Button("Tạo MP4")
            out_mp4 = gr.Video(label="Video MP4 dao động", format="mp4")

    out_text = gr.Textbox(label="Thông tin", interactive=False)

    # --- Ẩn/hiện ô ω khi nhập f ---
    def toggle_omega(f_val):
        if f_val is not None:
            return gr.update(visible=False)
        return gr.update(visible=True)

    f.change(toggle_omega, inputs=[f], outputs=[omega])

    btn_gif.click(generate_gif, inputs=[A,f,omega,phi,t_max], outputs=[out_text,out_gif])
    btn_mp4.click(generate_mp4, inputs=[A,f,omega,phi,t_max], outputs=[out_text,out_mp4])

if __name__ == "__main__":
    demo.launch()