import warnings 
from itertools import cycle

import numpy as np
import matplotlib.pyplot as plt


from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple

import numpy as np
import matplotlib.pyplot as plt


def _t_ppf(p: float, df: int) -> float:
    """
    Student-t inverse CDF (ppf). Only needed if CI is drawn.
    Uses SciPy if available; otherwise errors when CI is requested.
    """
    try:
        from scipy.stats import t as student_t  # type: ignore
    except Exception as e:
        raise ImportError(
            "Confidence intervals require SciPy (scipy.stats.t). "
            "Install with `pip install scipy`, or call draw(show_ci=False)."
        ) from e
    return float(student_t.ppf(p, df))


@dataclass
class _StepFit:
    x: float
    slope: float
    se_slope: float      # standard error of slope (sqrt(SSE) in MATLAB code)
    dfe: int             # degrees of freedom (n_pairs - 2)
    n_used: int          # number of error points used (after filtering/slicing)
    n_raw: int           # raw len(norms) for the step


class PlotConvergenceRate:
    """
    Convergence-rate plotter driven by xara.Model.testNorms().

    Assumes:
      - After each model.analyze(1), model.testNorms() returns the list
        of test norms for *that step's* Newton iterations.

    Computes (per step):
      log(e_{k+1}) = p * log(e_k) + c
    and stores p as the "convergence rate" (MATLAB slope).
    """

    def __init__(
        self,
        *,
        ci: Optional[float] = None,     # percent (e.g., 95). None/0 disables.
        n_ex_start: int = 1,            # ignore this many norms from start
        n_ex_end: int = 0,              # ignore this many norms from end
        n_total: int = 0,               # if >0, use only last n_total norms (>=3)
        step_start: float = 0.0,        # x-offset when x_mode="step"
        x_mode: str = "time",           # "step" or "time"
        x_label: Optional[str] = None,  # if None, chosen from x_mode
        y_label: str = "Convergence Rate",
        marker: str = None,
        markers = None,
        line_style: str = "none",       # MATLAB-like: markers only
        marker_size: float = 2.0,
        y_lim: Tuple[float, float] = (0.0, 5.0),
        y_ticks: Optional[Sequence[float]] = (0, 1, 2, 3, 4, 5),
        ci_alpha: float = 0.10,
        ax: Optional[plt.Axes] = None,
        skip: bool = False,
    ):
        self._skip = bool(skip)

        self.ci = None if (ci is None or ci == 0) else float(ci)
        self.n_ex_start = int(n_ex_start)
        self.n_ex_end = int(n_ex_end)
        self.n_total = int(n_total)
        self.step_start = float(step_start)
        self.x_mode = str(x_mode).strip().lower()

        if self.n_total != 0 and self.n_total < 3:
            raise ValueError("If n_total is used, it must be >= 3.")

        if self.x_mode not in ("step", "time"):
            raise ValueError('x_mode must be "step" or "time".')
        if markers is None and marker is None:
            markers = cycle(('.', 'o', 's', '^', 'v', 'x', '+'))
        elif marker is None:
            markers = cycle((marker,))

        self.markers = markers

        self.line_style = line_style
        self.marker_size = float(marker_size)
        self.y_lim = tuple(map(float, y_lim))
        self.y_ticks = None if y_ticks is None else list(map(float, y_ticks))
        self.ci_alpha = float(ci_alpha)

        # Figure/axes
        if not self._skip:
            if ax is None:
                fig, ax = plt.subplots()
                self.fig = fig
            else:
                self.fig = ax.figure
            self.ax = ax

            if x_label is None:
                x_label = "Load Step" if self.x_mode == "step" else "Load factor"
            self.ax.grid(True)
            self.ax.set_xlabel(x_label)
            self.ax.set_ylabel(y_label)
            self.ax.axhline(2, color="gray",  linestyle="-", linewidth=1)
            self.ax.axvline(0, color="black", linestyle="-", linewidth=1)
            self.ax.axhline(0, color="black", linestyle="-", linewidth=1)

        # Track global x-range over multiple series
        self._global_xmin = np.inf
        self._global_xmax = -np.inf

        self.reset(label="")

    def reset(self, label: str):
        """Start a new series (new line on the same axes)."""
        self._label = str(label)
        self._step_index = 0
        self._fits: List[_StepFit] = []
        self.marker = next(self.markers)

    def _step_x(self, model) -> float:
        """Choose x-coordinate for this step."""
        if self.x_mode == "time":
            # model.getTime() should exist in xara/opensees-style APIs
            return float(model.getTime())
        # "step" mode
        self._step_index += 1
        return self.step_start + self._step_index

    def _fit_slope_from_norms(self, norms: Sequence[float]) -> Tuple[float, float, int, int]:
        """
        Returns:
          slope, se_slope, dfe, n_used
        with NaNs / negative dfe if not computable.
        """
        e = np.asarray(list(norms), dtype=float).ravel()
        n_raw = int(e.size)

        if n_raw == 0:
            return np.nan, np.nan, -1, 0

        # Slice off start/end iterations (MATLAB: Punrm(1+NExStart:end-NExEnd))
        start = self.n_ex_start
        end = n_raw - self.n_ex_end if self.n_ex_end > 0 else n_raw
        if end <= start:
            return np.nan, np.nan, -1, 0

        e = e[start:end]

        # Ignore exactly 0 (MATLAB), and also guard against <=0 for log safety.
        e = e[np.isfinite(e)]
        e = e[e != 0.0]
        e = e[e > 0.0]

        n_used = int(e.size)

        # Need at least 3 error points to have >=2 pairs (MATLAB gate: NErr < 3 -> continue)
        if n_used < 3:
            return np.nan, np.nan, -1, n_used

        # If n_total is specified, use only the last n_total error points
        if self.n_total:
            if n_used < self.n_total:
                return np.nan, np.nan, -1, n_used
            e = e[-self.n_total :]
            n_used = int(e.size)

        # Regression pairs
        Xp = np.log(e[:-1])
        Yp = np.log(e[1:])

        devx = Xp - Xp.mean()
        devy = Yp - Yp.mean()
        sxx = float(np.sum(devx * devx))
        if sxx == 0.0:
            return np.nan, np.nan, -1, n_used

        slope = float(np.sum(devx * devy) / sxx)

        # Standard error of slope (for CI later)
        n_pairs = int(Xp.size)
        dfe = n_pairs - 2
        se_slope = np.nan
        if dfe > 0:
            intercept = float(Yp.mean() - slope * Xp.mean())
            resid = (Yp - intercept - slope * Xp)
            mse = float(np.sum(resid * resid) / dfe)
            sse = mse / sxx  # squared standard error of slope
            se_slope = float(np.sqrt(sse))

        return slope, se_slope, dfe, n_used

    def update(self, model):
        """
        Called once per load step (i.e., after each model.analyze(1) call).
        Reads norms from model.testNorms() and stores the fitted slope for that step.
        """
        if self._skip:
            return

        x = self._step_x(model)

        try:
            norms = model.testNorms()
        except Exception as e:
            raise AttributeError(
                "Model does not provide testNorms(), which is required for PlotConvergenceRate.update(model)."
            ) from e

        if norms is None:
            norms = []

        slope, se_slope, dfe, n_used = self._fit_slope_from_norms(norms)
        self._fits.append(
            _StepFit(
                x=x,
                slope=slope,
                se_slope=se_slope,
                dfe=dfe,
                n_used=n_used,
                n_raw=len(norms),
            )
        )

    def draw(self, *, show_ci: bool = False, **plot_kwargs):
        """
        Draw the current series.
        - Uses a new Matplotlib cycle color each time draw() is called.
        - If show_ci=True and ci is set, draws a translucent CI band.
        """
        if self._skip:
            return

        if not self._fits:
            warnings.warn("No steps recorded for this series.", stacklevel=2)
            return

        x = np.array([f.x for f in self._fits], dtype=float)
        y = np.array([f.slope for f in self._fits], dtype=float)

        have = np.isfinite(y)
        if not np.any(have):
            warnings.warn(
                "Not enough usable iterations in any step to determine convergence rate.",
                stacklevel=2,
            )
            return

        (line,) = self.ax.plot(
            x[have],
            y[have],
            label=self._label if self._label else None,
            linestyle=self.line_style,
            marker=self.marker,
            markersize=self.marker_size,
            **plot_kwargs,
        )
        color = line.get_color()

        # Optional CI (computed at draw time so SciPy is only required if you request CI)
        if show_ci and (self.ci is not None):
            lb = np.full_like(y, np.nan, dtype=float)
            ub = np.full_like(y, np.nan, dtype=float)

            alpha = 1.0 - self.ci / 100.0
            pcrit = 1.0 - alpha / 2.0

            for i, f in enumerate(self._fits):
                if not (np.isfinite(f.slope) and np.isfinite(f.se_slope) and f.dfe > 0):
                    continue
                t_score = _t_ppf(pcrit, f.dfe)
                me = t_score * f.se_slope
                lb[i] = f.slope - me
                ub[i] = f.slope + me

            have_ci = have & np.isfinite(lb) & np.isfinite(ub)
            if np.any(have_ci):
                self.ax.fill_between(
                    x[have_ci],
                    lb[have_ci],
                    ub[have_ci],
                    color=color,
                    alpha=self.ci_alpha,
                    linewidth=0,
                )

        # Update global x-range
        self._global_xmin = min(self._global_xmin, float(np.min(x[have])))
        self._global_xmax = max(self._global_xmax, float(np.max(x[have])))

    def finalize(self, *, title: Optional[str] = None):
        """Finalize formatting after all series are drawn."""
        if self._skip:
            return

        if title:
            self.ax.set_title(title)
        if len(self.ax.get_legend_handles_labels()):
            self.ax.legend()

        # MATLAB-ish y formatting
        self.ax.set_ylim(self.y_lim)
        if self.y_ticks is not None:
            self.ax.set_yticks(self.y_ticks)

        # Reasonable xlim across all series
        if np.isfinite(self._global_xmin) and np.isfinite(self._global_xmax):
            if self.x_mode == "step":
                # emulate MATLAB xlim([0 np] + StepStart) loosely
                self.ax.set_xlim(self.step_start, self._global_xmax + 1.0)
            else:
                self.ax.set_xlim(self._global_xmin, self._global_xmax)


    def savefig(self, filename: str, **savefig_kwargs):
        """Save the figure to a file."""
        if self._skip:
            return
        self.fig.savefig(filename, **savefig_kwargs)