"""
Solves a first-order ordinary differential equation (or system of first-order ODEs)
dy/dt = f(t, y) with initial condition y(t0) = y0.
"""
import numpy as np
from scipy.integrate import solve_ivp
from typing import Callable, List, Tuple, Dict, Any, Union
import matplotlib.pyplot as plt
from maths.differential_equations.ode_interface_utils import parse_float_list, parse_time_span, string_to_ode_func
import gradio as gr

ODEFunc = Callable[[float, Union[np.ndarray, List[float]]], Union[np.ndarray, List[float]]]

def solve_first_order_ode(
    ode_func: ODEFunc,
    t_span: Tuple[float, float],
    y0: List[float],
    t_eval_count: int = 100,
    method: str = 'RK45',
    **kwargs: Any
) -> Dict[str, Union[np.ndarray, str, bool]]:
    # ...existing code...
    try:
        y0_np = np.array(y0, dtype=float)
        t_eval = np.linspace(t_span[0], t_span[1], t_eval_count)

        sol = solve_ivp(ode_func, t_span, y0_np, method=method, t_eval=t_eval, **kwargs)

        plot_path = None
        if sol.success:
            try:
                plt.figure(figsize=(10, 6))
                if y0_np.ndim == 0 or len(y0_np) == 1 : # Single equation
                    plt.plot(sol.t, sol.y[0], label=f'y(t), y0={y0_np[0] if y0_np.ndim > 0 else y0_np}')
                else: # System of equations
                    for i in range(sol.y.shape[0]):
                        plt.plot(sol.t, sol.y[i], label=f'y_{i+1}(t), y0_{i+1}={y0_np[i]}')
                plt.xlabel("Time (t)")
                plt.ylabel("Solution y(t)")
                plt.title(f"Solution of First-Order ODE ({method})")
                plt.legend()
                plt.grid(True)
                plot_path = "ode_solution_plot.png"
                plt.savefig(plot_path)
                plt.close() # Close the plot to free memory
            except Exception as e_plot:
                print(f"Warning: Could not generate plot: {e_plot}")
                plot_path = None

        return {
            't': sol.t,
            'y': sol.y,
            'message': sol.message,
            'success': sol.success,
            'plot_path': plot_path
        }
    except Exception as e:
        return {
            't': np.array([]),
            'y': np.array([]),
            'message': f"Error during ODE solving: {str(e)}",
            'success': False,
            'plot_path': None
        }
    # --- Gradio Interface for First-Order ODEs ---
first_order_ode_interface = gr.Interface(
    fn=lambda ode_str, t_span_str, y0_str, t_eval_count, method: solve_first_order_ode(
        string_to_ode_func(ode_str, ('t', 'y')),
        parse_time_span(t_span_str),
        parse_float_list(y0_str),
        int(t_eval_count),
        method
    ),
    inputs=[
        gr.Textbox(label="ODE Function (lambda t, y: ...)",
                   placeholder="e.g., lambda t, y: -y*t  OR for system lambda t, y: [y[1], -0.1*y[1] - y[0]]",
                   info="Define dy/dt or a system [dy1/dt, dy2/dt,...]. `y` is a list/array for systems."),
        gr.Textbox(label="Time Span (t_start, t_end)", placeholder="e.g., 0,10"),
        gr.Textbox(label="Initial Condition(s) y(t_start)", placeholder="e.g., 1 OR for system 1,0"),
        gr.Slider(minimum=10, maximum=1000, value=100, step=10, label="Evaluation Points Count"),
        gr.Radio(choices=['RK45', 'LSODA', 'BDF', 'RK23', 'DOP853'], value='RK45', label="Solver Method")
    ],
    outputs=[
        gr.Image(label="Solution Plot", type="filepath", show_label=True, visible=lambda res: res['success'] and res['plot_path'] is not None),
        gr.Textbox(label="Solver Message"),
        gr.Textbox(label="Success Status"),
        gr.JSON(label="Raw Data (t, y values)", visible=lambda res: res['success']) # For users to copy if needed
    ],
    title="First-Order ODE Solver",
    description="""
Solves dy/dt = f(t, y) or a system of first-order ODEs.

- Enter a Python lambda for the ODE (e.g., `lambda t, y: -y*t`).
- For systems, use `y` as a list: `lambda t, y: [y[1], -0.1*y[1] - y[0]]`.
- Initial conditions: single value (e.g., `1`) or comma-separated for systems (e.g., `1,0`).

**Examples:**

- Simple: `lambda t, y: -y*t`, y0: `1`, t_span: `0,5`
- System: `lambda t, y: [y[1], -0.1*y[1] - y[0]]`, y0: `1,0`, t_span: `0,20`
- Lotka-Volterra: `lambda t, y: [1.5*y[0] - 0.8*y[0]*y[1], 0.5*y[0]*y[1] - 0.9*y[1]]`, y0: `10,5`, t_span: `0,20`

WARNING: Uses eval() for the ODE function string - potential security risk.
""",
    flagging_mode="manual"
)