"""Parameters for the simulation."""

from argparse import Action
from dataclasses import asdict, dataclass, fields
from typing import Any, Literal, NotRequired, TypedDict

from . import solvers

# Type definitions for parameter names
ElementType = Literal["Cobalt", "Iron", "Nickel"]
SolverType = Literal["opencl", "mpi", "numpy", "jax"]
PrecisionType = Literal["single", "double"]
DeviceType = Literal["cpu", "gpu", "auto"]
InitType = Literal["0", "dw"]


@dataclass
class RunParameters:
    """
    Store simulation parameters.

    Example:
        >>> params = RunParameters(element="Cobalt", N=1000)
        >>> print(params)
        element         : Cobalt
        N               = 1000
        dt              = 1e-14
        Jx              = 300
        Jy              = 21
        Jz              = 21
        dx              = 1e-09
        T               = 1100
        H_ext           = 0.0
        init_type       : 0
        result_file     : run.npz
        start_averaging = 2000
        n_mean          = 1
        n_profile       = 0
        solver          : numpy
        precision       : double
        blocking        = False
        seed            = 12345
        device          : auto
        np              = 1
    """

    element: ElementType = "Cobalt"  #: Chemical element of the sample
    N: int = 5000  #: Number of iterations
    dt: float = 1.0e-14  #: Time step
    Jx: int = 300  #: Number of points in x
    Jy: int = 21  #: Number of points in y
    Jz: int = 21  #: Number of points in z
    dx: float = 1.0e-9  #: Step in x
    T: float = 0.0  #: Temperature (K)
    H_ext: float = 0.0  #: External field (A/m)
    init_type: InitType = "0"  #: Initialization type
    result_file: str = "run.npz"  #: Result file name
    start_averaging: int = 4000  #: Start index of time average
    n_mean: int = 1  #: Spatial average frequency (number of iterations)
    n_profile: int = 0  #: x-profile save frequency (number of iterations)
    solver: SolverType = "numpy" if solvers.size == 1 else "mpi"  #: Solver to use
    precision: PrecisionType = "double"  #: Precision of the simulation
    blocking: bool = False  #: Use blocking communications
    seed: int = 12345  #: Random seed for temperature fluctuations
    device: DeviceType = "auto"
    profiling: bool = False  #: Enable profiling output
    np: int = solvers.size  #: Number of processes (for MPI solver)

    def as_dict(self) -> dict[str, Any]:
        """
        Convert RunParameters to a dictionary.

        Returns:
            Dictionary representation of the parameters
        """
        return asdict(self)

    def __str__(self) -> str:
        """Return solver information in a readable format."""
        # Compute the width of the longest parameter name
        width = max([len(field.name) for field in fields(self)])
        s = ""
        for field in fields(self):
            value = getattr(self, field.name)
            # the seprator is ":" for strings and "=" for others
            sep = ":" if isinstance(value, str) else "="
            s += "{0:<{1}} {2} {3}\n".format(field.name, width, sep, value)
        return s


DEFAULT_RUN_PARAMETERS = RunParameters()
DRP = DEFAULT_RUN_PARAMETERS  #: Alias for default parameters

ParameterValue = str | int | float | bool


class ArgParameter(TypedDict):
    """Metadata for a simulation parameter."""

    help: str
    default: ParameterValue
    choices: NotRequired[Any]
    action: NotRequired[str | type[Action]]
    type: NotRequired[type]


def lit_to_list(lit: Any) -> list[Any]:
    """
    Convert a Literal type to a list of its possible values.

    Args:
        lit: Literal type

    Returns:
        List of possible values
    """
    return list(lit.__args__)


arg_parameters: dict[str, ArgParameter] = {
    "element": {
        "help": "Chemical element of the sample",
        "default": DRP.element,
        "choices": lit_to_list(ElementType),
    },
    "N": {"help": "Number of time iterations", "default": DRP.N},
    "dt": {"help": "Time step", "default": DRP.dt},
    "Jx": {"help": "Number of points in x", "default": DRP.Jx},
    "Jy": {"help": "Number of points in y", "default": DRP.Jy},
    "Jz": {"help": "Number of points in z", "default": DRP.Jz},
    "dx": {"help": "Step in x", "default": DRP.dx},
    "T": {"help": "Temperature (K)", "default": DRP.T},
    "H_ext": {"help": "External field (A/m)", "default": DRP.H_ext},
    "init_type": {
        "help": "Type of initialization ('0' for uniform, 'dw' for domain wall)",
        "default": DRP.init_type,
        "choices": lit_to_list(InitType),
    },
    "result_file": {
        "help": "Name of the npz result file",
        "default": DRP.result_file,
    },
    "start_averaging": {
        "help": "Start index of time average",
        "default": DRP.start_averaging,
    },
    "n_mean": {
        "help": "Spatial average frequency (number of iterations)",
        "default": DRP.n_mean,
    },
    "n_profile": {
        "help": "x-profile save frequency (number of iterations)",
        "default": DRP.n_profile,
    },
    "solver": {
        "help": "Solver to use for the simulation",
        "default": DRP.solver,
        "choices": lit_to_list(SolverType),
    },
    "precision": {
        "help": "Precision of the floating point (single or double)",
        "default": DRP.precision,
        "choices": lit_to_list(PrecisionType),
    },
    "blocking": {
        "help": "Use blocking communications (MPI solver only)",
        "default": DRP.blocking,
        "action": "store_true",
    },
    "seed": {
        "help": "Random seed for temperature fluctuations",
        "default": DRP.seed,
        "type": int,
    },
    "device": {
        "help": "Device to use by the OpenCL solver",
        "default": DRP.device,
        "choices": lit_to_list(DeviceType),
    },
    "profiling": {
        "help": "Enable profiling output (internal profiler)",
        "default": DRP.profiling,
        "action": "store_true",
    },
}  #: simulation CLI parameters
