"""Axes Actions such as xaxis, noaxais, etc."""

import re
from typing import Literal

from finesse.analysis.actions.base import Action
from finesse.parameter import Parameter, ParameterRef

from .sweep import Sweep, get_sweep_array


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class Noxaxis(Sweep):
    """The default Action for a model when no action is specified. Solves the model
    without changing any parameters. Will return a solution with a single value
    for every detector in the model.

    Parameters
    ----------
    pre_step : Action | None, optional
        An action to perform before the Noxaxis is run, by default None
    post_step : Action | None, optional
        An action to perform after the Noxaxis is run, by default None
    name : str, optional
        Name of the action, used to find the solution in the final output, by
        default "noxaxis"
    """

    def __init__(
        self,
        pre_step: Action | None = None,
        post_step: Action | None = None,
        name: str = "noxaxis",
    ):
        super().__init__(name=name, pre_step=pre_step, post_step=post_step)


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class XNaxis(Sweep):
    """N-dimensional version of Xaxis allows for multi-dimensional parameter sweep
    of arbitrary dimension N. You need to provide a set of 'parameter, mode, start,
    stop, steps' arguments for every dimension you want to sweep over.

    Parameters
    ----------
    relative : bool, optional
        Whether to apply the changes to the parameter value on top of the current
        value or ignore the current value of the parameter, by default False, by
        default False
    pre_step : Action | None, optional
        An action to perform before each step is computed, by default None
    post_step : Action | None, optional
        An action to perform after each step is computed, by default None
    name : str, optional
        Name of the action, used to find the solution in the final output, by
        default "XNaxis"
    """

    def __init__(
        self,
        *args,
        relative: bool = False,
        pre_step: Action | None = None,
        post_step: Action | None = None,
        name: str = "XNaxis",
    ):
        if len(args) % 5 != 0:
            raise Exception(
                f"XNaxis arguments must come in groups of five: parameter, mode, start, stop, steps not {args}"
            )
        self.relative = relative
        self.N = len(args) // 5

        if self.N == 0:
            raise Exception("XNaxis requires at least one axis to be specified")
        # Here we map the XNaxis arguments to the Sweep inputs
        self.__set_args = args
        new_args = []

        for i in range(0, len(args), 5):
            new_args.append(args[i + 0])
            new_args.append(
                get_sweep_array(args[i + 2], args[i + 3], args[i + 4], args[i + 1])
            )
            new_args.append(relative)
        super().__init__(*new_args, pre_step=pre_step, post_step=post_step, name=name)

    def __getattr__(self, key):
        res = re.match("(parameter|mode|start|stop|steps)([0-9]*)", key)
        if res is None:
            super().__getattribute__(key)
        else:
            grp = res.groups()
            N = 1 if grp[1] == "" else int(grp[1])
            if N == 0:
                raise Exception("Specify an axes greater than 0")
            if N > self.N:
                raise Exception(f"This xaxis does not have {N} axes")
            idx = 5 * (N - 1)
            if grp[0] == "parameter":
                return self.__set_args[idx + 0]
            elif grp[0] == "mode":
                return self.__set_args[idx + 1]
            elif grp[0] == "start":
                return self.__set_args[idx + 2]
            elif grp[0] == "stop":
                return self.__set_args[idx + 3]
            elif grp[0] == "steps":
                return self.__set_args[idx + 4]


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class Xaxis(XNaxis):
    """Sweeps a parameter and solves the model for every value in the sweep. Will
    a solution with an array of values for every detector in the model.

    Parameters
    ----------
    parameter : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start : float
        Starting value of the parameter sweep
    stop : float
        Final value of the parameter sweep.
    steps : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    relative : bool, optional
        Whether to apply the changes to the parameter value on top of the current
        value or ignore the current value of the parameter, by default False
    pre_step : Action | None, optional
        An action to perform before each step is computed, by default None
    post_step : Action | None, optional
        An action to perform after each step is computed, by default None
    name : str, optional
        Name of the action, used to find the solution in the final output, by
        default "xaxis"
    """

    def __init__(
        self,
        parameter: Parameter | ParameterRef | str,
        mode: Literal["lin", "log"],
        start: float,
        stop: float,
        steps: int,
        relative: bool = False,
        *,
        pre_step: Action | None = None,
        post_step: Action | None = None,
        name: str = "xaxis",
    ):
        super().__init__(
            parameter,
            mode,
            start,
            stop,
            steps,
            relative=relative,
            pre_step=pre_step,
            post_step=post_step,
            name=name,
        )


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class X2axis(XNaxis):
    """Two-dimensional version of the Xaxis action. Will solve the model for every
    point on a grid of values defined by the two sets of sweep arguments. Returns a
    solution with a two-dimensional array for every detector in the model.

    Significantly faster than calling Xaxis in a Python for loop.

    Parameters
    ----------
    parameter1 : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode1 : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start1 : float
        Starting value of the parameter sweep
    stop1 : float
        Final value of the parameter sweep.
    steps1 : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    parameter2 : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode2 : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start2 : float
        Starting value of the parameter sweep
    stop2 : float
        Final value of the parameter sweep.
    steps2 : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    relative : bool, optional
        Whether to apply the changes to the parameter value on top of the current
        value or ignore the current value of the parameter, by default False, by
        default False
    pre_step : Action | None, optional
        An action to perform before each step is computed, by default None
    post_step : Action | None, optional
        An action to perform after each step is computed, by default None
    name : str, optional
        Name of the action, used to find the solution in the final output, by
        default "x2axis"
    """

    def __init__(
        self,
        parameter1: Parameter | ParameterRef | str,
        mode1: Literal["lin", "log"],
        start1: float,
        stop1: float,
        steps1: int,
        parameter2: Parameter | ParameterRef | str,
        mode2: Literal["lin", "log"],
        start2: float,
        stop2: float,
        steps2: int,
        relative: bool = False,
        *,
        pre_step: Action | None = None,
        post_step: Action | None = None,
        name: str = "x2axis",
    ):
        super().__init__(
            parameter1,
            mode1,
            start1,
            stop1,
            steps1,
            parameter2,
            mode2,
            start2,
            stop2,
            steps2,
            relative=relative,
            pre_step=pre_step,
            post_step=post_step,
            name=name,
        )


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class X3axis(XNaxis):
    """Three-dimensional version of the Xaxis action. Will solve the model for every
    point on a cube of values defined by the three sets of sweep arguments. Returns
    a solution with a three-dimensional array for every detector in the model.

    Significantly faster than calling Xaxis in a nested Python for loop.

    Parameters
    ----------
    parameter1 : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode1 : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start1 : float
        Starting value of the parameter sweep
    stop1 : float
        Final value of the parameter sweep.
    steps1 : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    parameter2 : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode2 : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start2 : float
        Starting value of the parameter sweep
    stop2 : float
        Final value of the parameter sweep.
    steps2 : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    parameter3 : Parameter | ParameterRef | str
        Parameter to sweep over.
    mode3 : Literal['lin', 'log']
        Whether to sweep linearly or logarithmically
    start3 : float
        Starting value of the parameter sweep
    stop3 : float
        Final value of the parameter sweep.
    steps3 : int
        Number of steps in the sweep. Solution arrays will contain
        'steps+1' values
    relative : bool, optional
        Whether to apply the changes to the parameter value on top of the current
        value or ignore the current value of the parameter, by default False, by
        default False
    pre_step : Action | None, optional
        An action to perform before each step is computed, by default None
    post_step : Action | None, optional
        An action to perform after each step is computed, by default None
    name : str, optional
        Name of the action, used to find the solution in the final output, by
        default "x3axis"
    """

    def __init__(
        self,
        parameter1: Parameter | ParameterRef | str,
        mode1: Literal["lin", "log"],
        start1: float,
        stop1: float,
        steps1: int,
        parameter2: Parameter | ParameterRef | str,
        mode2: Literal["lin", "log"],
        start2: float,
        stop2: float,
        steps2: int,
        parameter3: Parameter | ParameterRef | str,
        mode3: Literal["lin", "log"],
        start3: float,
        stop3: float,
        steps3: int,
        relative: bool = False,
        *,
        pre_step: Action | None = None,
        post_step: Action | None = None,
        name="x3axis",
    ):
        super().__init__(
            parameter1,
            mode1,
            start1,
            stop1,
            steps1,
            parameter2,
            mode2,
            start2,
            stop2,
            steps2,
            parameter3,
            mode3,
            start3,
            stop3,
            steps3,
            relative=relative,
            pre_step=pre_step,
            post_step=post_step,
            name=name,
        )
