import abc
import sys
import warnings
from datetime import datetime
from time import ctime
from typing import Callable
from typing import Optional

import numpy as np

from .evolve import EvolutionModel
from .progress import DummyProgressBar
from .progress import get_pbar
from dustyn._typing import FloatLike
from dustyn.core.record import Record


class SolverRecord(Record):
    def __init__(self, *args, **kwargs):

        print(
            "warning: the SolverRecord class is now identical to Record (except for this warning).",
            file=sys.stderr,
        )
        super().__init__(*args, **kwargs)


class Solver(abc.ABC):
    @staticmethod
    @abc.abstractmethod
    def step(
        R: np.ndarray, time: FloatLike, timestep: FloatLike, func: Callable
    ) -> None:
        """This should walk the R vector state over one timestep according to the evolution
        function `func` with signature

        def func(time, R):
            ...
        """
        pass

    @staticmethod
    def _update_from_buffer(output: list[np.ndarray], buffer=np.ndarray) -> None:
        output.append(buffer.copy())

    @classmethod
    def integrate(
        cls,
        model: EvolutionModel,
        R0: np.ndarray,
        *,
        tstart: FloatLike,
        tstop: FloatLike,
        CFL: FloatLike,
        pbar: bool = False,
        buffer_size: int = 1000,
        dt_record: Optional[FloatLike] = None,
        dt_min: FloatLike = 0.0,
        label: Optional[str] = None,
        logging: bool = False,
    ) -> Record:

        if CFL > 1.0:
            warnings.warn(
                f"Received {CFL=} > 1, stability is not guaranteed and results will likely be inaccurate."
            )
        elif CFL < 0:
            raise ValueError("Received negative CFL security.")
        elif CFL < (min_float := np.finfo(R0.dtype).eps):
            raise ValueError(
                f"Received too low value for {CFL=} . The minimal value is {min_float}"
            )

        tstart, tstop, CFL = map(R0.dtype.type, (tstart, tstop, CFL))

        if pbar:
            if buffer_size < 1000:
                warnings.warn(
                    "refreshing the progress bar too often is extremely expensive. "
                    "It is strongly advised to use progress bars only with `buffer_size` >= 1000."
                )
            if (
                dt_record is not None
                and int((tstop - tstart) // dt_record) <= 2 * buffer_size
            ):
                warnings.warn(
                    "Time resolution is too low to allow useful granulosity in the progressbar."
                )
            if logging:
                warnings.warn(
                    "logging and pbar can not be used simultaneously. Turning logging off."
                )
                logging = False

            title = label if label is not None else "Integrating..."
            progress = get_pbar(title=title, total=tstop)
        else:
            progress = DummyProgressBar()

        metadata = dict(
            CFL=CFL,
            label=label,
            geometry=getattr(model, "geometry", None),
            dimensionality=getattr(model, "dimensionality", None),
            stop_reason="completed",
        )
        states_buffer = np.empty((buffer_size, *R0.shape), dtype=R0.dtype)
        times_buffer = np.empty(buffer_size, dtype=R0.dtype)

        buffer_pos = 0
        R = R0.copy()
        t = tstart

        times: list[np.ndarray] = [np.array([tstart])]
        states: list[np.ndarray] = [np.expand_dims(R0, axis=0)]
        rec_count = 1
        next_trec = dt_record
        timestep_phy = CFL * model.get_max_timestep(R)
        it = 0
        while t < tstop:
            if dt_record is None:
                timestep_rec = min(timestep_phy, tstop - t)
                next_trec = t + timestep_rec
            else:
                timestep_rec = next_trec - t
            to_rec: bool = timestep_rec <= timestep_phy
            timestep = timestep_rec if to_rec else timestep_phy

            cls.step(R, time=t, timestep=timestep, func=model.evolve)

            t += timestep

            # this is essentially equivalent to
            # >>> timestep = min(CFL * model.get_max_timestep(R), tsave - t)
            # but with the advantage of being floating-point comparison proof
            # when we need to decide wether or not to save later.
            timestep_phy = CFL * model.get_max_timestep(R)
            nan_timestep: bool = timestep_phy != timestep_phy

            if to_rec:
                # t is updated manually here to secure floating point
                # equality check in the buffer updating block
                np.testing.assert_approx_equal(t, next_trec, significant=14)
                t = next_trec
                states_buffer[buffer_pos] = R
                times_buffer[buffer_pos] = t
                buffer_pos += 1
                rec_count += 1
                if dt_record is not None:
                    next_trec = min(dt_record * rec_count, tstop)

                if logging:
                    print(
                        f"{ctime(datetime.now().timestamp())} | {t=} | {it=}",
                        file=sys.stderr,
                    )
            elif nan_timestep or timestep < dt_min:
                if logging:
                    warnings.warn(
                        f"Aborting run with {timestep=} (minvalue is {dt_min=})"
                    )
                states_buffer[buffer_pos] = R
                times_buffer[buffer_pos] = t
                buffer_pos += 1
                rec_count += 1
                if logging:
                    print(
                        f"{ctime(datetime.now().timestamp())} | {t=} | {it=}",
                        file=sys.stderr,
                    )

            if (
                buffer_pos == buffer_size
                or t == tstop
                or nan_timestep
                or timestep < dt_min
            ):
                cls._update_from_buffer(states, states_buffer[:buffer_pos])
                cls._update_from_buffer(times, times_buffer[:buffer_pos])
                progress.update(completed=t, refresh=True)
                buffer_pos = 0

            it += 1

            if nan_timestep:
                metadata["stop_reason"] = "NaN timestep"
                break
            elif timestep < dt_min:
                metadata["stop_reason"] = "reached minimal timestep"
                break

        progress.update(completed=tstop, refresh=True)
        progress.stop()

        return Record(
            times=np.concatenate(times),
            states=np.concatenate(states),
            metadata=metadata,
        )


class RK4(Solver):
    """
    This is the Runge-Kutta method, with order 4
    https://en.wikipedia.org/wiki/Runge–Kutta_methods
    """

    @staticmethod
    def step(
        R: np.ndarray, time: FloatLike, timestep: FloatLike, func: Callable
    ) -> None:
        # - I'm avoiding premature optimisation, so I'm keeping the details of time positioning (t+dt)
        #   even though it should make no difference in my applications
        k1 = func(time, R)
        k2 = func(time + timestep / 2, R + k1 * timestep / 2)
        k3 = func(time + timestep / 2, R + k2 * timestep / 2)
        k4 = func(time + timestep, R + k3 * timestep)
        R += timestep / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
