import gradio as gr
import numpy as np
from control.matlab import tf, pole, zero, frequency_response, bode
from sympy.polys.polytools import poly_from_expr
import plotly.graph_objects as go
from plotly.subplots import make_subplots


def poly_handler(poly_str: str) -> list[float]:
    try:
        if poly_str.isdigit():
            poly_coffs = [poly_str]
        else:
            poly = poly_from_expr(poly_str)[0]
            poly_coffs = poly.all_coeffs()

        return list(map(float, poly_coffs))

    except Exception as e:
        raise ValueError(f"Error parsing polynomial: {e}")


def bode_information(poly_upper: str, poly_lower: str) -> tuple:
    upper_poly_coffs = poly_handler(poly_upper)
    lower_poly_coffs = poly_handler(poly_lower)

    num, den = np.array(upper_poly_coffs), np.array(lower_poly_coffs)

    try:
        trans_func_g = tf(num, den)
    except Exception as e:
        raise ValueError(f"Error creating transfer function: {e}")

    pole_result = pole(trans_func_g)
    zero_result = zero(trans_func_g)
    mag, phase, omega = frequency_response(trans_func_g)

    return mag, phase, omega


def to_latex(poly_upper: str, poly_lower: str) -> str:
    upper_poly_coffs = poly_handler(poly_upper)
    lower_poly_coffs = poly_handler(poly_lower)

    def process_item(i: float, c: float):
        c = f"+ {c}" if c >= 0 else f"- {-c}"

        if i == 0:
            return f"{c}"
        if i == 1:
            return f"{c} s"

        return f"{c} s^{i}"

    upper_str = "".join(
        [process_item(i, c) for i, c in enumerate(upper_poly_coffs[::-1]) if c != 0]
    )
    lower_str = "".join(
        [process_item(i, c) for i, c in enumerate(lower_poly_coffs[::-1]) if c != 0]
    )

    if upper_str.startswith("+ "):
        upper_str = upper_str[2:]
    if lower_str.startswith("+ "):
        lower_str = lower_str[2:]

    latex_str = r"\frac{" + upper_str + r"}{" + lower_str + r"}"

    return latex_str


def bode_graph(mag: np.ndarray, phase: np.ndarray, omega: np.ndarray):
    mag_db = 20 * np.log10(mag)
    fig = make_subplots(
        rows=1,
        cols=2,
        subplot_titles=(
            "Bode Plot - Magnitude",
            "Bode Plot - Phase",
        ),
    )

    fig.add_trace(
        go.Scatter(x=omega, y=mag_db, mode="lines", name="Magnitude"), row=1, col=1
    )
    fig.add_trace(
        go.Scatter(x=omega, y=phase, mode="lines", name="Phase"), row=1, col=2
    )

    fig.update_yaxes(title_text="Magnitude (dB)", row=1, col=1)
    fig.update_yaxes(title_text="Phase (degrees)", row=1, col=2)
    fig.update_xaxes(type="log", title_text="Frequency (rad/s)", row=1, col=1)
    fig.update_xaxes(type="log", title_text="Frequency (rad/s)", row=1, col=2)
    fig.update_layout(title_text="Bode Plot")

    return fig


def run_it(upper_poly: str, lower_poly: str):
    mag, phase, omega = bode_information(upper_poly, lower_poly)
    latex_str = to_latex(upper_poly, lower_poly)
    fig = bode_graph(mag, phase, omega)

    return fig, f"$$H(s)={latex_str}$$"


def main():
    with gr.Blocks() as demo:
        gr.Markdown("# Bode Plotter")
        gr.Markdown("This app allows you to plot Bode plots for transfer functions.")
        gr.Markdown(
            "Enter the transfer function in the form of a string, e.g., '1/(s^2 + 2*s + 1)'."
        )
        upper_input = gr.Textbox(
            label="Numerator Polynomial",
            placeholder="Enter the numerator polynomial (e.g., '1 - 2*s + 3*s^2')",
            lines=2,
        )
        lower_input = gr.Textbox(
            label="Denominator Polynomial",
            placeholder="Enter the denominator polynomial (e.g., '1 + 4*s + 5*s^2')",
            lines=2,
        )
        run_btn = gr.Button("Plot Bode", variant="primary")
        gr.Markdown("### Bode Plot")

        with gr.Column():
            latex_output = gr.Markdown(label="Transfer Function LaTeX")
            bode_plot = gr.Plot(label="Bode Plot")

        run_btn.click(
            run_it,
            inputs=[
                upper_input,
                lower_input,
            ],
            outputs=[bode_plot, latex_output],
        )

    demo.launch()


if __name__ == "__main__":
    main()