#! /usr/bin/env python
import os
import pathlib
import textwrap
import warnings
from collections import defaultdict
from functools import partial
from io import StringIO
from typing import Optional, TextIO

import numpy as np
import rich_click as click
import tomlkit as tomllib
import yaml
from landlab import RasterModelGrid
from landlab.io.netcdf import write_raster_netcdf
from packaging.version import parse as parse_version
from scipy import interpolate

from ._version import __version__
from .plume import Plume

click.rich_click.ERRORS_SUGGESTION = (
    "Try running the '--help' flag for more information."
)
click.rich_click.ERRORS_EPILOGUE = (
    "To find out more, visit https://github.com/csdms/bmi-wavewatch3"
)
click.rich_click.STYLE_ERRORS_SUGGESTION = "yellow italic"
click.rich_click.SHOW_ARGUMENTS = True
click.rich_click.GROUP_ARGUMENTS_OPTIONS = False
click.rich_click.SHOW_METAVARS_COLUMN = True
click.rich_click.USE_MARKDOWN = True
out = partial(click.secho, bold=True, err=True)
err = partial(click.secho, fg="red", err=True)


# class RiverParameters:
#     name = "river"
#     width = Parameter(
#         "width", default=50.0, valid=Range(0.0, max=None), help="river width", units="m",
#     )
#     filepath = Parameter(
#         "filepath", default="river.csv", valid=Path(exists=True, file_okay=True, dir_okay=False), help="river file",
#     )
#
#
# class SedimentParameters:
#     name = "sediment"
#     removal_rate = Parameter(
#     )
#     bulk_density = Parameter(
#     )
#
#
# class OceanParameters:
#     name = "ocean"
#     filepath = Parameter(
#         "filepath", default="ocean.csv", valid=CsvFile(exists=True, nrows="+", ncols=3), help="river file",
#     )
#
#
# class GridParameters:
#     name = "grid"
#     shape = Parameter("shape", [500, 500], valid=Array(length=2, dtype=int), help="number of grid rows and columns")
#     shape = Parameter("shape", [500, 500], valid=Length(2), help="number of grid rows and columns")
#     xy_spacing = Parameter("xy_spacing", [100.0, 100.0], valid=Length(2), help="spacing of grid columns and rows")
#     xy_of_lower_left = Parameter("xy_of_lower_left", [0.0, 0.0], valid=Length(2), help="coordinates of lower-left node of grid")


def load_config(file: Optional[TextIO] = None):
    """Load plume config file.

    Parameters
    ----------
    fname : file-like, optional
        Opened config file or ``None``. If ``None``, return default
        values.

    Returns
    -------
    dict
        Config parameters.
    """
    conf = {
        # "_version": __version__,
        "grid": {
            "shape": [500, 500],
            "xy_spacing": [100.0, 100.0],
            "xy_of_lower_left": [0.0, 0.0],
        },
        "river": {
            # "filepath": "river.csv",
            "width": 50.0,
            "depth": 5.0,
            "velocity": 1.5,
            "location": [0.0, 25000.0],
            "angle": 0.0,
        },
        "sediment": {"removal_rate": 60.0, "bulk_density": 1600.0},
        "ocean": {
            # "filepath": "ocean.csv",
            "along_shore_velocity": 0.1,
            "sediment_concentration": 0.0,
        },
        "output": {"filepath": "plume.nc"},
    }
    if file is not None:
        conf.update(yaml.safe_load(file))

    conf.setdefault("_version", __version__)

    this_version = parse_version(__version__)
    that_version = parse_version(conf["_version"])
    if this_version.major != that_version.major:
        warnings.warn(
            f"possible version mismatch. file is version v{that_version},"
            f" but you are using plume v{this_version}",
            stacklevel=2,
        )

    return conf


def _contents_of_input_file(infile: str) -> str:
    params = load_config()

    def as_csv(data, header=None):
        with StringIO() as fp:
            np.savetxt(fp, data, header=header, delimiter=",", fmt="%.1f")
            contents = fp.getvalue()
        return contents

    contents = {
        "plume.toml": tomllib.dumps({"plume": params}),
        # "river.csv": as_csv(
        #     [[0.0, 50.0, 5.0, 1.5]],
        #     header=os.linesep.join(
        #         [
        #             f"version: {__version__}",
        #             "Time [d], Width [m], Depth [m], Velocity [m/s]",
        #         ]
        #     ),
        # ),
        # "ocean.csv": as_csv(
        #     [[0.0, 0.1, 0.0]],
        #     header=os.linesep.join(
        #         [
        #             f"version: {__version__}",
        #             "Time [d], Along-shore velocity [m/s], Sediment Concentration [-]",
        #         ]
        #     ),
        # ),
    }

    return contents[infile]


def load_params_from_strings(values):
    params = defaultdict(dict)

    for param in values:
        group_dot_name, value = param.split("=")
        value = yaml.load(value)
        try:
            group, name = group_dot_name.split(".")
        except ValueError:
            name = group_dot_name
            params[name] = value
        else:
            params[group][name] = value

    return params


@click.group(chain=True)
@click.version_option(package_name="sed-plume")
@click.option(
    "--cd",
    default=".",
    type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True),
    help="chage to directory, then execute",
)
@click.option(
    "-s",
    "--silent",
    is_flag=True,
    help="Suppress status status messages, including the progress bar.",
)
@click.option(
    "-v", "--verbose", is_flag=True, help="Also emit status messages to stderr."
)
def plume(cd: str, silent: bool, verbose: bool) -> None:
    """Simulate a hypopycnal plume.

    ## Examples

    Create a folder with example input files,

    ```bash
    $ plume setup example
    ```

    Run a simulation using the examples input files,

    ```bash
    $ plume run example
    ```
    """
    if not silent and verbose:
        out(f"run_dir = {cd!r}")
    os.chdir(cd)


class RiverTimeSeries:
    def __init__(self, filepath, kind="linear"):
        data = np.loadtxt(filepath, delimiter=",", comments="#").reshape((-1, 4))
        if len(data) == 1:
            data = np.vstack([data, data])
            data[1, 0] = data[0, 0] + 1

        self._river = interpolate.interp1d(
            data[:, 0],
            data[:, 1:],
            kind=kind,
            axis=0,
            copy=True,
            assume_sorted=True,
            bounds_error=False,
            fill_value=(data[0, 1:], data[-1, 1:]),
        )
        self._time = 0.0
        self._width, self._depth, self._velocity = self._river(self._time)

    @property
    def velocity(self):
        return self._velocity

    @property
    def width(self):
        return self._width

    @property
    def depth(self):
        return self._depth

    def update(self):
        self._time += 1.0
        self._width, self._depth, self._velocity = self._river(self._time)


class OceanTimeSeries:
    def __init__(self, filepath, kind="linear"):
        data = np.loadtxt(filepath, delimiter=",", comments="#").reshape((-1, 3))
        if len(data) == 1:
            data = np.vstack([data, data])
            data[1, 0] = data[0, 0] + 1

        self._ocean = interpolate.interp1d(
            data[:, 0],
            data[:, 1:],
            kind=kind,
            copy=True,
            assume_sorted=True,
            bounds_error=False,
            fill_value=(data[0, 1:], data[-1, 1:]),
        )
        self._time = 0.0
        self._velocity, _ = self._ocean(self._time)

    @property
    def velocity(self):
        return self._velocity

    def update(self):
        self._time += 1.0
        self._velocity, _ = self._ocean(self._time)


@plume.command()
@click.option("--dry-run", is_flag=True, help="Do not actually run the model")
@click.pass_context
def run(ctx, dry_run: bool) -> None:
    """Run the plume simulation on a set of input files.

    ## Examples

    ```bash
    $ mkdir example && cd example
    $ plume setup
    $ plume run
    """
    verbose = ctx.parent.params["verbose"]
    silent = ctx.parent.params["silent"]

    # params = load_params("plume.toml")
    with open("plume.toml") as fp:
        params = tomllib.load(fp)

    if verbose and not silent:
        out("plume.toml = |")
        out(textwrap.indent(tomllib.dumps({"plume": params}).strip(), prefix="  "))

    params = params["plume"]

    if dry_run:
        out("Nothing to do. 😴")
    else:
        output_file = pathlib.Path(params["output"]["filepath"])
        if output_file.exists():
            err(f"{output_file}: output files exists")
            raise click.Abort()

        grid = RasterModelGrid(
            params["grid"]["shape"],
            xy_spacing=params["grid"]["xy_spacing"],
            xy_of_lower_left=params["grid"]["xy_of_lower_left"],
        )

        grid.add_zeros("sediment_deposit_thickenss", at="node")  # .reshape(grid.shape)
        grid.at_grid["sediment__removal_rate"] = params["sediment"]["removal_rate"]
        grid.at_grid["sediment__bulk_density"] = params["sediment"]["bulk_density"]

        # river = RiverTimeSeries(params["river"]["filepath"])
        # ocean = OceanTimeSeries(params["ocean"]["filepath"])

        params["river"]["angle"] = np.deg2rad(params["river"]["angle"])

        # n_days = 2
        # for day in range(n_days):

        plume = Plume(
            grid,
            river_width=params["river"]["width"],
            river_depth=params["river"]["depth"],
            river_velocity=params["river"]["velocity"],
            river_angle=params["river"]["angle"],
            river_loc=params["river"]["location"],
            ocean_velocity=params["ocean"]["along_shore_velocity"],
        )
        plume.run_one_step()

        write_raster_netcdf(
            output_file,
            plume.grid,
            time=0.0,
            append=True,
            attrs={"_version": __version__},
        )

        out("💥 Finished! 💥")
        out(f"Output written to {output_file}")

        print(output_file)


@plume.command()
@click.argument("infile", type=click.Choice(sorted(["plume.toml"])))
def generate(infile: str) -> None:
    """Print example input files.

    ## Examples

    To see an example *ocean.csv* input file,
    ```bash
    $ plume generate ocean.csv
    ```
    """
    print(_contents_of_input_file(infile))


@plume.command()
@click.pass_context
def setup(ctx) -> None:
    """Setup a folder of input files for a simulation.

    *plume setup* creates a set of example input files for the *plume* program.
    This set of input files can then be used with the *plume run* command.

    > **NOTE**: *plume setup* will not overwrite any of your files. If you want to overwrite
    existing files, you will have to remove them yourself and then re-run
    *plume setup*.

    ## Examples

    ```bash
    $ mkdir _example
    $ plume --cd=_example setup
    ```
    """
    verbose = ctx.parent.params["verbose"]
    silent = ctx.parent.params["silent"]

    folder = pathlib.Path(".")

    files = [pathlib.Path(fname) for fname in ["plume.toml"]]

    existing_files = [folder / name for name in files if (folder / name).exists()]
    if existing_files:
        for name in existing_files:
            err(
                f"{name}: File exists. Either remove and then rerun or choose a different destination folder",
            )
    else:
        for fname in files:
            with open(folder / fname, "w") as fp:
                print(_contents_of_input_file(str(fname)), file=fp)
        if verbose and not silent:
            out(f"files = {[str(f) for f in files]!r}")
        print(str(folder.absolute()))

    if existing_files:
        raise click.Abort()
