import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from sympy import latex

def newtons_method(f, x, a,
                  # --- Newton Config ---
                  iterations=5, tolerance=1e-10,
                  stoppingcriterion="absolute",
                  output="value",
                  # --- Visibility Toggles ---
                  showfunction=True, showpoints=True,
                  showroot=False, showtangents=True, showverticallines=True,
                  showinfo=False, showlegend=True,
                  # --- Labels & Text ---
                  title="", caption="",
                  xlabel="", ylabel="",
                  # --- Styling ---
                  colors=None, line_styles=None, linewidth=2, fontsize=None,
                  # --- Style Options (dicts) ---
                  functionoptions=None, pointoptions=None,
                  tangentoptions=None, verticallineoptions=None, rootoptions=None,
                  # --- Limits ---
                  xlim=None, ylim=None, view=None,
                  # --- Animation Config ---
                  animation_time=5000,
                  # --- Figure Size ---
                  width=800, height=500):

    # Label Formatting
    def _smart_label(lbl):
        if lbl is None:
            return None
        if isinstance(lbl, str):
            if any(c in lbl for c in ['_', '^', '\\']):
                return f"${lbl}$"
            return lbl
        return f"${latex(lbl)}$"

    # Display LaTeX summary line
    def _display_latex_summary(seq):
        from IPython.display import Math, display as ipy_display
        n_iters = len(seq) - 1
        f_latex  = latex(f)
        ipy_display(Math(
            rf"{n_iters} \text{{ iterations of Newton's method applied to }}"
            rf"f(x) = {f_latex}"
            rf"\text{{ with initial point }} p_0 = {float(a):.4f}"
        ))

    # Font Scaling
    def _get_font_sizes():
        if fontsize is None:
            scale    = min(width / 800, height / 500)
            title_fs = max(4, 16 * scale)
            label_fs = max(4, 14 * scale)
            tick_fs  = max(4, 12 * scale)
            leg_fs   = max(4, 12 * scale)
        else:
            title_fs = fontsize + 2
            label_fs = fontsize
            tick_fs  = max(4, fontsize - 2)
            leg_fs   = fontsize
        return title_fs, label_fs, tick_fs, leg_fs

    # Resolve view / xlim / ylim
    def _resolve_limits(seq, f_num):
        nonlocal xlim, ylim

        if view is not None:
            if isinstance(view[0], (int, float)):
                xlim = (float(view[0]), float(view[1]))
            else:
                xlim = (float(view[0][0]), float(view[0][1]))
                ylim = (float(view[1][0]), float(view[1][1]))

        if xlim is None:
            x_span = max(abs(max(seq) - min(seq)) * 2.0, 2.0)
            x_cen  = np.mean(seq)
            xlim   = (x_cen - x_span, x_cen + x_span)

        x_plot = np.linspace(xlim[0], xlim[1], 500)
        y_plot = f_num(x_plot)
        finite = y_plot[np.isfinite(y_plot)]

        if ylim is None:
            if len(finite):
                pad  = (finite.max() - finite.min()) * 0.15 or 1.0
                ylim = (finite.min() - pad, finite.max() + pad)
            else:
                ylim = (-5, 5)

        return xlim, ylim, x_plot, y_plot

    # Build default style dicts
    def _build_styles():
        _c = colors      or []
        _l = line_styles or []

        func_opts  = dict(color=_c[0] if len(_c) > 0 else '#0055A4', 
                          linewidth=linewidth,
                          linestyle=_l[0] if len(_l) > 0 else '-',
                          zorder=3)
        tan_opts   = dict(color=_c[1] if len(_c) > 1 else '#E55523', 
                          linewidth=1.2,
                          linestyle=_l[1] if len(_l) > 1 else '-',
                          zorder=4)
        vert_opts  = dict(color=_c[2] if len(_c) > 2 else '#E55523',
                          linewidth=1.2,
                          linestyle=_l[2] if len(_l) > 2 else ':',   
                          zorder=4)
        point_opts = dict(color=_c[3] if len(_c) > 3 else '#4CAF50', 
                          marker='o', markersize=5,
                          markerfacecolor='none',                    
                          linestyle='None', zorder=5)
        root_opts  = dict(color=_c[4] if len(_c) > 4 else 'red',
                          marker='s', markersize=7,
                          linestyle='None', zorder=6)

        if functionoptions:      func_opts.update(functionoptions)
        if tangentoptions:       tan_opts.update(tangentoptions)
        if verticallineoptions:  vert_opts.update(verticallineoptions)
        if pointoptions:         point_opts.update(pointoptions)
        if rootoptions:          root_opts.update(rootoptions)

        return func_opts, tan_opts, vert_opts, point_opts, root_opts

    # Run Newton iterations
    def _run_iterations(f_num, dfdx_num):
        _valid = ("absolute", "relative", "function", "residual")
        if stoppingcriterion not in _valid:
            raise ValueError(f"Invalid stoppingcriterion='{stoppingcriterion}'. "
                             f"Choose from: {_valid}")

        xi  = float(a)
        seq = [xi]

        for _ in range(iterations):
            fi  = float(f_num(xi))
            dfi = float(dfdx_num(xi))

            if abs(dfi) < 1e-14:
                print("Warning: f'(x) ≈ 0, stopping.")
                break

            xi_new = xi - fi / dfi

            if stoppingcriterion in ("absolute", "function", "residual"):
                converged = abs(fi) < tolerance
            elif stoppingcriterion == "relative":
                converged = abs(xi_new - xi) < tolerance * (1.0 + abs(xi))

            seq.append(xi_new)
            xi = xi_new

            if converged:
                break

        return seq

    # Setup Figure & Axes
    def _setup_figure(xlim, ylim, auto_title, auto_caption):
        title_fs, label_fs, tick_fs, _ = _get_font_sizes()

        fig, ax = plt.subplots(figsize=(width / 100, height / 100))
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        
        ax.grid(True, linestyle='-', alpha=0.3, zorder=0)

        # Center axes like Maple
        ax.spines['left'].set_position('zero')
        ax.spines['bottom'].set_position('zero')
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')
        
        # Move labels to edges (only applies if string is not empty)
        ax.set_xlabel(_smart_label(xlabel), fontsize=label_fs, loc='right')
        ax.set_ylabel(_smart_label(ylabel), fontsize=label_fs, loc='top')

        _title = _smart_label(title) if title is not None else auto_title
        if _title:
            ax.set_title(_title, fontsize=title_fs, pad=20)

        ax.tick_params(axis='both', labelsize=tick_fs)

        _cap = auto_caption if caption is None else caption
        if _cap:
            ax.text(0.5, -0.09, _cap, transform=ax.transAxes,
                    ha='center', fontsize=max(8, title_fs - 4), color='gray')

        return fig, ax

    # Draw steps for static plot
    def _draw_steps(ax, seq, k, f_num, dfdx_num,
                    tan_opts, vert_opts, point_opts, root_opts, xlim):
        
        if showpoints and k >= 0:
            ax.plot([seq[0]], [0], **point_opts) # Initial guess point

        for i in range(min(k, len(seq) - 1)):
            xi_  = seq[i]
            xi1_ = seq[i + 1]
            fi_  = float(f_num(xi_))

            if showverticallines:
                ax.plot([xi_, xi_], [0, fi_], **vert_opts)

            if showtangents:
                ax.plot([xi_, xi1_], [fi_, 0], **tan_opts)

            if showpoints:
                ax.plot([xi_],  [fi_],  **point_opts)
                ax.plot([xi1_], [0],    **point_opts)

        if showroot and len(seq) > 1:
            ax.plot([seq[-1]], [0], **root_opts, label=f"root ≈ {seq[-1]:.8f}")

    # plot
    def _make_plot(seq, f_num, dfdx_num):
        func_opts, tan_opts, vert_opts, point_opts, root_opts = _build_styles()
        xlim_, ylim_, x_plot, y_plot = _resolve_limits(seq, f_num)
        _, _, _, leg_fs = _get_font_sizes()

        auto_title   = f"Newton's Method:  $f(x) = {latex(f)}$,   $x_0 = {a}$"
        auto_caption = (f"Converged to x ≈ {seq[-1]:.8f}  after {len(seq)-1} iteration(s)"
                        f"  [{stoppingcriterion} criterion, tol={tolerance}]")

        fig, ax = _setup_figure(xlim_, ylim_, auto_title, auto_caption)

        if showfunction:
            ax.plot(x_plot, y_plot, label=f"$f(x) = {latex(f)}$", **func_opts)

        _draw_steps(ax, seq, len(seq) - 1, f_num, dfdx_num,
                    tan_opts, vert_opts, point_opts, root_opts, xlim_)

        if showlegend:
            handles, labels = ax.get_legend_handles_labels()
            if handles:
                ax.legend(fontsize=leg_fs, loc="lower left")

        plt.tight_layout()
        plt.show()
        _display_latex_summary(seq)
        return None

    # animation
    def _make_animation(seq, f_num, dfdx_num):
        plt.rcParams["animation.html"]        = "jshtml"
        plt.rcParams['animation.embed_limit'] = 50.0
        plt.ioff()

        func_opts, tan_opts, vert_opts, point_opts, root_opts = _build_styles()
        xlim_, ylim_, x_plot, y_plot = _resolve_limits(seq, f_num)
        _, _, _, leg_fs = _get_font_sizes()

        auto_title   = f"Newton's Method:  $f(x) = {latex(f)}$,   $x_0 = {a}$"
        auto_caption = (f"Converged to x ≈ {seq[-1]:.8f}  after {len(seq)-1} iteration(s)"
                        f"  [{stoppingcriterion} criterion, tol={tolerance}]")

        fig, ax = _setup_figure(xlim_, ylim_, auto_title, auto_caption)

        if showfunction:
            ax.plot(x_plot, y_plot, label=f"$f(x) = {latex(f)}$", **func_opts)
            
        if showlegend:
            handles, labels = ax.get_legend_handles_labels()
            if handles:
                ax.legend(fontsize=leg_fs, loc="lower left")

        n = len(seq)
        tan_lines  = [ax.plot([], [], **tan_opts)[0]   for _ in range(n - 1)]
        vert_lines = [ax.plot([], [], **vert_opts)[0]  for _ in range(n - 1)]
        curve_pts  = [ax.plot([], [], **point_opts)[0] for _ in range(n)]
        xaxis_pts  = [ax.plot([], [], **point_opts)[0] for _ in range(n - 1)]
        root_pt,   = ax.plot([], [], **root_opts)
        
        start_pt,  = ax.plot([seq[0]], [0], **point_opts)
        start_pt.set_visible(showpoints)

        iter_text = ax.text(0.02, 0.96, '', transform=ax.transAxes,
                            fontsize=max(8, leg_fs), va='top',
                            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
                            visible=showinfo)

        all_artists = tan_lines + vert_lines + curve_pts + xaxis_pts + [root_pt, start_pt, iter_text]

        def _update_frame(frame):
            k = frame

            for i in range(n - 1):
                if i < k:
                    xi_  = seq[i];  xi1_ = seq[i + 1]
                    fi_  = float(f_num(xi_))

                    if showverticallines:
                        vert_lines[i].set_data([xi_, xi_], [0, fi_])
                    if showtangents:
                        tan_lines[i].set_data([xi_, xi1_], [fi_, 0])
                    if showpoints:
                        curve_pts[i].set_data([xi_],  [fi_])
                        xaxis_pts[i].set_data([xi1_], [0])
                else:
                    tan_lines[i].set_data([], [])
                    vert_lines[i].set_data([], [])
                    curve_pts[i].set_data([], [])
                    xaxis_pts[i].set_data([], [])

            if showpoints and k < n:
                curve_pts[k].set_data([seq[k]], [float(f_num(seq[k]))])

            if showroot and k == n - 1:
                root_pt.set_data([seq[-1]], [0])
            else:
                root_pt.set_data([], [])

            if showinfo:
                fi_display = abs(float(f_num(seq[min(k, n - 1)])))
                iter_text.set_text(
                    f"Step {k} / {n - 1}\n"
                    f"x      = {seq[min(k, n-1)]:.8f}\n"
                    f"|f(x)| = {fi_display:.2e}\n"
                    f"criterion: {stoppingcriterion}"
                )

            return all_artists

        interval = animation_time / max(n, 1)
        anim = FuncAnimation(fig, _update_frame, frames=n,
                             interval=interval, repeat=True, blit=True)

        plt.close(fig)

        from IPython.display import HTML, display as ipy_display
        ipy_display(HTML(anim.to_jshtml()))
        _display_latex_summary(seq)
        return anim

    # Main Execution
    dfdx     = sp.diff(f, x)
    f_num    = sp.lambdify(x, f,    modules='numpy')
    dfdx_num = sp.lambdify(x, dfdx, modules='numpy')

    seq = _run_iterations(f_num, dfdx_num)

    if output == "value": return seq[-1]
    if output == "sequence": return seq
    if output == "plot": return _make_plot(seq, f_num, dfdx_num)
    if output == "animation": return _make_animation(seq, f_num, dfdx_num)

    raise ValueError(f"Invalid output='{output}'. Choose: 'value', 'sequence', 'plot', 'animation'.")