"""
Command Line Interface
----------------------
"""

from __future__ import annotations

import json
import re
from importlib.util import find_spec
from pathlib import Path
from typing import Any, TextIO
from warnings import warn

import click
import numpy as np
import yaml

import pigreads.plot as pp
from pigreads import Models, __version__
from pigreads.gui import LiveView, get_iz_iv
from pigreads.progress import PROGRESS_ITERS
from pigreads.schema.model import ModelDefinition
from pigreads.schema.simulation import Simulation

_cmc_available: bool = False
if find_spec("cmcrameri") is not None:
    import cmcrameri  # type: ignore[import-untyped]

    assert cmcrameri
    _cmc_available = True


@click.group()
@click.option(
    "-m",
    "--define-model",
    "model_defs",
    type=click.File("r"),
    multiple=True,
    help="Define a model using a YAML file. "
    "The beginning of the filename will be used as the model key.",
)
def main(model_defs: list[TextIO]) -> None:
    "Run and visualise Pigreads simulations."

    for model_def in model_defs:
        stem = Path(model_def.name).stem
        match = re.match(r"^[0-9A-Za-z_]+", stem)
        message = f"Filename '{stem}' must start with an alphanumeric identifier!"
        assert match, message
        key = match.group(0)
        message = f"Defining Models.available['{key}'] using {model_def.name}..."
        warn(message, RuntimeWarning, stacklevel=2)
        Models.available[key] = ModelDefinition(**yaml.safe_load(model_def))


@main.command()
def version() -> None:
    "Show installed version."
    click.echo(f"Pigreads {__version__}")


@main.command()
@click.argument("config", type=click.File("r"))
@click.argument("result", type=click.Path())
@click.option(
    "-p",
    "--progress",
    type=click.Choice(list(PROGRESS_ITERS.keys())),
    default="bar",
    show_default=True,
    help="Progress bar type.",
)
@click.option(
    "-s", "--start-frame", type=int, default=0, help="Index of frame to start with."
)
@click.option(
    "-l",
    "--live",
    type=yaml.safe_load,
    default="{}",
    show_default=True,
    help="YAML encoded keyword arguments for LiveView to show a live preview of "
    "a chosen variable. Instead of a YAML string, a path to a YAML file may "
    "also be given, or the name or index of a variable.",
)
@click.option(
    "-i",
    "--infinite",
    is_flag=True,
    default=False,
    help="Run the simulation in an infinite loop.",
)
def run(
    config: TextIO,
    result: click.Path,
    progress: str,
    start_frame: int,
    live: int | str | dict[str, Any],
    infinite: bool,
) -> None:
    """
    Run a simulation.

    CONFIG is a YAML or JSON file defining the simulation,
    and RESULT is a path to write the results to.
    """

    sim = Simulation(**yaml.safe_load(config))
    live_: dict[str, Any] = LiveView.normalise_kwargs(live)
    if len(live_) > 0 and "cmap" not in live_:
        live_["cmap"] = "cmc.acton" if _cmc_available else None
    liveview = (
        LiveView(models=sim.prepare_models(), dy=sim.dy, dx=sim.dx, **live_)
        if len(live_) > 0
        else None
    )

    def callback(states: np.ndarray[Any, Any], ifr: int) -> bool:
        return (
            True
            if liveview is None
            else liveview.update(states[ifr], time=sim.Nt * sim.dt * ifr)
        )

    sim.run(
        path=str(result),
        progress=PROGRESS_ITERS[progress],
        start_frame=start_frame,
        callback=callback,
        infinite=infinite,
    )


@main.command()
@click.argument("input", type=click.File("r"), default="-")
@click.argument("output", type=click.File("w"), default="-")
@click.option(
    "-f",
    "--format",
    type=click.Choice(["yaml", "json"]),
    default="yaml",
    show_default=True,
    help="Output format.",
)
def convert(
    input: TextIO,  # pylint: disable=redefined-builtin
    output: TextIO,
    format: str,  # pylint: disable=redefined-builtin
) -> None:
    """
    Convert a configuration file to canonical format.

    INPUT is a YAML or JSON file defining the simulation,
    and OUTPUT is a path to write the converted file to.
    """
    sim = Simulation(**yaml.safe_load(input))
    data = sim.model_dump()

    if format == "yaml":
        yaml.safe_dump(data, output, allow_unicode=True)
    else:
        assert format == "json"
        json.dump(data, output, indent=2)


@main.command()
@click.argument("config", type=click.File("r"))
@click.argument("result", type=click.Path(exists=True))
@click.argument("movie", type=click.Path())
@click.option(
    "-d", "--dpi", type=int, default=180, show_default=True, help="Dots per inch."
)
@click.option(
    "-r", "--fps", type=int, default=15, show_default=True, help="Frames per second."
)
@click.option(
    "-p",
    "--progress",
    type=click.Choice(list(PROGRESS_ITERS.keys())),
    default="bar",
    show_default=True,
    help="Progress bar type.",
)
@click.option("-z", "--index-z", type=int, default=None, help="Index in z direction.")
@click.option(
    "-v", "--variable", type=str, default=None, help="Variable name or index."
)
@click.option(
    "--kwargs",
    type=yaml.safe_load,
    default="{}",
    show_default=True,
    help="Additional arguments as YAML or JSON to pass to movie, "
    "plot_frame, and imshow, for example "
    '\'{"cmap": "inferno", "vmin": -0.25, "vmax": 1.25}\'.',
)
@click.option(
    "-n",
    "--parallel",
    type=int,
    default=1,
    show_default=True,
    help="Number of processes. 0 to use all CPUs.",
)
def movie(
    config: TextIO,
    result: click.Path,
    movie: click.Path,  # pylint: disable=redefined-outer-name
    dpi: int,
    fps: int,
    progress: str,
    variable: str | None,
    index_z: int | None,
    parallel: int,
    kwargs: dict[str, Any],
) -> None:
    """
    Generate a movie from a result file using FFmpeg and Matplotlib.

    CONFIG is a YAML or JSON file defining the simulation,
    RESULT is a NPY file containing the simulation results,
    and MOVIE is a path to write an MP4 video to.
    """

    sim = Simulation(**yaml.safe_load(config))
    models = sim.prepare_models()

    array: np.ndarray[Any, Any] = np.lib.format.open_memmap(
        str(result),
        mode="r",
    )
    assert array.shape == (sim.Nfr, sim.Nz, sim.Ny, sim.Nx, models.Nv)
    iz, iv = get_iz_iv(models, variable, index_z or sim.Nz // 2)
    array = array[:, iz, :, :, iv]

    zlabel = f" z = {sim.dz * iz:.2f}" if sim.Nz > 1 else ""
    tlables = [
        f"t = {time}" + zlabel for time in (sim.dt * sim.Nt * np.arange(sim.Nfr))
    ]

    pp.movie(
        path=str(movie),
        frames=array,
        tlables=tlables,
        progress=progress,
        parallel=parallel,
        **pp.imshow_defaults(
            sim=sim,
            array=array,
            dpi=dpi,
            fps=fps,
            vlabel=list(models.available[models[0].key].variables)[iv],
            **kwargs,
        ),
    )


__all__ = [
    "convert",
    "main",
    "movie",
    "run",
    "version",
]
