#! /usr/bin/env python3

"""Utility functions for testing scan planners.

These including fake sample creation, scan path visualisation, and
persistent storage of expected scan paths for samples.
"""

import argparse
import os
import pickle

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from matplotlib.patches import PathPatch
from matplotlib.path import Path as MatPath
from scipy import interpolate

from openflexure_microscope_server import scan_planners

THIS_DIR = os.path.dirname(os.path.realpath(__file__))

ALL_SAMPLE_NAMES = ("lobed", "regular", "core")


class FakeSample:
    """A fake sample to test scan algorithms.

    The sample is able to return whether a given position is sample, there is no image
    associated with the sample
    """

    def __init__(self, xy_points: list[tuple[int, int]]):
        """Create the sample from a spline interpolation around the given points."""
        self._sample_perimeter = interp_closed_path(xy_points, 500)

    def is_sample(self, pos: tuple[int, int], im_size: tuple[int, int]) -> bool:
        """Return True if an image at a given location is on the sample.

        The image size is specified to check if it overlaps the sample. It doesn't
        check the entire image field as this is designed to be used where the fake
        sample is much larger than the image and has smooth edges. It just checks the
        4 corners.
        """
        img_corners = [
            (pos[0] + im_size[0], pos[1] + im_size[1]),
            (pos[0] + im_size[0], pos[1] - im_size[1]),
            (pos[0] - im_size[0], pos[1] + im_size[1]),
            (pos[0] - im_size[0], pos[1] - im_size[1]),
        ]
        return any(
            self._sample_perimeter.contains_point(corner) for corner in img_corners
        )

    @property
    def patch(self) -> PathPatch:
        """The sample as a matplotlib patch for plotting."""
        patch = PathPatch(self._sample_perimeter)
        patch.set(color=(1.0, 0.8, 1.0, 1.0))
        return patch


def visualise_scan(sample: FakeSample, planner: scan_planners.ScanPlanner) -> Figure:
    """For a given sample and scanner object return a matplotlib figure of the scan."""
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.add_artist(sample.patch)
    xh, yh = zip(*planner.path_history, strict=True)
    xi, yi, _zi = zip(*planner.imaged_locations, strict=True)
    if planner.secondary_locations:
        xs, ys, _zs = zip(*planner.secondary_locations, strict=True)
    else:
        xs, ys, _zs = [], [], []

    # convert history to numpy array so can calculate quiver arrows
    xh = np.array(xh)
    yh = np.array(yh)

    plt.quiver(
        xh[:-1],
        yh[:-1],
        xh[1:] - xh[:-1],
        yh[1:] - yh[:-1],
        scale_units="xy",
        angles="xy",
        scale=1,
    )
    plt.plot(xh, yh, "r.")
    plt.plot(xi, yi, "g*")
    plt.plot(xs, ys, "o", mfc="none", mec="blue")
    ax.axis("equal")
    return fig


def interp_closed_path(xy_points: list[tuple[int, int]], n_points: int) -> MatPath:
    """Interpolate an n_point closed curve from a lists of xy_points.

    This can be used to create an arbitrary sample shape for testing a scan planner.

    Modified from:
    https://stackoverflow.com/questions/33962717/interpolating-a-closed-curve-using-scipy
    """
    # Use zip to separate x and y points into tuples
    x, y = zip(*xy_points, strict=True)

    # Append first point and convert to array
    x = np.array(x + (x[0],))
    y = np.array(y + (y[0],))

    # fit splines to x=f(u) and y=g(u), treating both as periodic. also note that s=0
    # is needed in order to force the spline fit to pass through all the input points.
    spline_data, *_unused = interpolate.splprep([x, y], s=0, per=True)

    # evaluate the spline
    xi, yi = interpolate.splev(np.linspace(0, 1, n_points), spline_data)

    # Convert to a matplotlib closed path
    path_points = [[xp, yp] for xp, yp in (zip(xi, yi, strict=True))]
    return MatPath(path_points, closed=True)


def example_smart_spiral(
    sample_name: str = "lobed",
) -> tuple[FakeSample, scan_planners.ScanPlanner]:
    """Run an example scan.

    :returns: The sample scanned and the planner object after scan is complete.
    """
    xy_sample_points = load_sample_points(sample_name)
    sample = FakeSample(xy_sample_points)
    img_size = (1000, 1000)
    initial_position = (0, 0)
    planner_settings = {"dx": 1200, "dy": 800, "max_dist": 100000}
    planner = scan_planners.SmartSpiral(
        initial_position=initial_position, planner_settings=planner_settings
    )

    while not planner.scan_complete:
        xy_pos, _ = planner.get_next_location_and_z_estimate()
        xyz_pos = (xy_pos[0], xy_pos[1], 0)
        imaged = sample.is_sample(xy_pos, img_size)
        planner.mark_location_visited(xyz_pos, imaged=imaged, focused=imaged)
    return sample, planner


def profile_example_smart_spiral():
    """Profile running an example scan and print the cumulative profile stats."""
    import cProfile
    import pstats

    profiler = cProfile.Profile()
    profiler.runcall(example_smart_spiral)

    stats_fname = os.path.join(THIS_DIR, "scan_example_stats.pstats")
    profiler.dump_stats(stats_fname)

    run_stats = pstats.Stats(stats_fname)
    run_stats.strip_dirs()
    run_stats.sort_stats("cumulative")
    run_stats.print_stats("scan_planners.py")


def plot_all_examples():
    """Plot all examples as png files."""
    for sample_name in ALL_SAMPLE_NAMES:
        sample, planner = example_smart_spiral(sample_name)
        png_fname = os.path.join(THIS_DIR, f"scan_example_plot_{sample_name}.png")
        fig = visualise_scan(sample, planner)
        fig.savefig(png_fname, dpi=200)


def update_example_smart_spiral_pickles():
    """Pickle the ScanPlanner for the example_smart_spiral().

    This is done so the history can be compared by testing to check
    the algorithm is unchanged.

    If the algorithm is purposefully changed then this will need to be
    run to update the pickle for the test to pass.

    Takes sample, the sample type we have generated, so we can make a
    pickle for each sample type
    """
    for sample_name in ALL_SAMPLE_NAMES:
        pkl_fname = os.path.join(THIS_DIR, f"example_smart_spiral_{sample_name}.pkl")
        _, planner = example_smart_spiral(sample_name)
        with open(pkl_fname, "wb") as pkl_file_obj:
            pickle.dump(planner, pkl_file_obj, pickle.HIGHEST_PROTOCOL)


def get_expected_result_for_example_smart_spiral(
    sample_name: str,
) -> scan_planners.ScanPlanner:
    """Return the expected ScanPlanner object for the example_smart_spiral().

    This is loaded from a pickle so that the object can be committed to the repo.
    """
    pkl_fname = os.path.join(THIS_DIR, f"example_smart_spiral_{sample_name}.pkl")
    with open(pkl_fname, "rb") as pkl_file_obj:
        return pickle.load(pkl_file_obj)


def load_sample_points(sample_name: str):
    """Return the points to generate the FakeSample corresponding to the given input name.

    Options are "lobed", "regular", and "core".
    """
    sample_options = {
        "lobed": [
            (-5000, -5000),
            (-2000, 16000),
            (1000, 2000),
            (6000, 7000),
            (9000, 2000),
        ],
        "regular": [
            (-5000, -5000),
            (-5000, 5000),
            (5000, 5000),
            (5000, -5000),
        ],
        "core": [
            (-12000, 2000),
            (-12000, 1000),
            (0, -2000),
            (10000, -2000),
            (10000, -1000),
            (1000, 0),
        ],
    }
    if sample_name not in sample_options:
        all_samples = ", ".join(sample_options.keys())
        raise ValueError(
            f"{sample_name} is not a valid sample name. Valid names : {all_samples}"
        )
    return sample_options[sample_name]


def main():
    """Run the profiler, the plotting, or update the pickles based on command line input.

    This only runs if run as a command line script.
    """
    parser = argparse.ArgumentParser(description="Simulated scan-planning utility")

    subparsers = parser.add_subparsers(dest="command", required=True)

    # profile
    subparsers.add_parser("profile", help="Run simulation under cProfile")

    # plot
    subparsers.add_parser("plot", help="Run simulation and produce plots")

    # update
    subparsers.add_parser("update", help="Update stored reference pickles")

    args = parser.parse_args()

    if args.command == "profile":
        profile_example_smart_spiral()
    elif args.command == "plot":
        plot_all_examples()
    elif args.command == "update":
        update_example_smart_spiral_pickles()
    else:
        parser.error(f"Unknown command: {args.command}")


if __name__ == "__main__":
    main()
