import time
from datetime import timedelta
from math import sqrt
from pathlib import Path
from statistics import fmean
from time import sleep

import linuxpy.video.device
import requests
from linuxpy.video.device import Device, PixelFormat, VideoCapture

ESP3D_HOST = "10.177.42.107"

MAX_POS = (250, 232)
CENTER_POS = (MAX_POS[0] / 2, MAX_POS[1] / 2)

ITEM_SIZE = (250, 92)

CAMERA_FRAME_WIDTH = 1920
CAMERA_FRAME_HEIGHT = 1080
CAMERA_FRAME_PIXEL_FORMAT = PixelFormat.MJPEG
CAMERA_FRAME_SKIP = 5

class Stage:
    def __init__(self, host: str, accel_decel_time_margin: float = 0.3):
        self.command_url = f"http://{host}/command"
        self.x: float | None = None
        self.y: float | None = None
        self.accel_decel_time_margin = accel_decel_time_margin

    def wait_seconds_for_motion(self, x: float, y: float, speed: float):
        distance = sqrt(x * x + y * y)
        return (distance * 100) / speed + self.accel_decel_time_margin

    def wait_for_motion(self, x: float, y: float, speed: float):
        sleep(self.wait_seconds_for_motion(x, y, speed))

    def move_relative(self, x: float, y: float, speed: float = 1000):
        if self.x is None or self.y is None:
            raise ValueError("Unknown absolute position")
        if self.x + x < 0 or self.x + x > MAX_POS[0] or \
            self.y + y < 0 or self.y + y > MAX_POS[1]:
            raise ValueError(f"Target position ({x}, {y}) outside of bounds (0-{MAX_POS[0]}, 0-{MAX_POS[1]})")
        requests.get(self.command_url, params=(("cmd", f"G91\nG1 X{x} Y{y} F{speed}"),))
        self.wait_for_motion(x, y, speed)
        self.x += x
        self.y += y

    def move_absolute(self, x: float, y: float, speed: float = 1000, initial: bool = False):
        if x < 0 or x > MAX_POS[0] or y < 0 or y > MAX_POS[1]:
            raise ValueError(f"Target position ({x}, {y}) outside of bounds (0-{MAX_POS[0]}, 0-{MAX_POS[1]})")
        if (self.x is None or self.y is None) and not initial:
            raise ValueError("Unknown absolute position and not an initial move")
        requests.get(self.command_url, params=(("cmd", f"G90\nG1 X{x} Y{y} F{speed}"),))
        if not initial:
            self.wait_for_motion(x - self.x, y - self.y, speed)
        self.x = x
        self.y = y

def open_camera() -> Device:
    for device in linuxpy.video.device.iter_devices():
        device.open()
        if 'YW' in device.info.card and any(map(lambda f: f.pixel_format == CAMERA_FRAME_PIXEL_FORMAT, device.info.formats)):
            return device
        device.close()
    raise FileNotFoundError

def snake_path(start_x: int, start_y: int, end_x: int, end_y: int):
    step_x = 4
    step_y = 4
    for x in range(start_x, end_x + step_x, step_x):
        for y in range(start_y, end_y + step_y, step_y) if x // step_x % 2 == 0 else range(end_y, start_y - step_y, -step_y):
            yield x, y

def main():
    cam = open_camera()
    stage = Stage(ESP3D_HOST)

    output_directory = Path("/tmp/microscope")
    output_directory.mkdir(parents=True, exist_ok=True)
    for entry in output_directory.iterdir():
        try:
            entry.unlink(missing_ok=True)
        except IsADirectoryError:
            pass

    start_x = CENTER_POS[0] - ITEM_SIZE[0] / 2
    end_x = CENTER_POS[0] + ITEM_SIZE[0] / 2
    start_y = CENTER_POS[1] - ITEM_SIZE[1] / 2
    end_y = CENTER_POS[1] + ITEM_SIZE[1] / 2

    with (output_directory / "path.svg").open("w") as svg:
        svg.write(f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {MAX_POS[0]} {MAX_POS[1]}">')
        svg.write(f'<path fill="none" stroke-width="0.2" stroke="red" d="')
        for i, (x, y) in enumerate(snake_path(int(start_x), int(start_y), int(end_x), int(end_y))):
            svg.write(f'{"M" if i == 0 else "L"}{x} {y} ')
        svg.write('"/>')
        for x, y in snake_path(int(start_x), int(start_y), int(end_x), int(end_y)):
            svg.write(f'<circle cx="{x}" cy="{y}" r="0.4" fill="#700"/>')
        svg.write('</svg>')

    inset = 10
    stage.move_absolute(end_x - inset, start_y + inset, initial=True)
    input("Adjust focus for corner 1, press enter when done")
    stage.move_absolute(end_x - inset, end_y - inset)
    input("Adjust focus for corner 2, press enter when done")
    stage.move_absolute(start_x + inset, end_y - inset)
    input("Adjust focus for corner 3, press enter when done")
    stage.move_absolute(start_x + inset, start_y + inset)
    input("Adjust focus for corner 4, press enter when done")
    stage.move_absolute(start_x, start_y) #, initial=True)
    #input("Press enter when in position")

    path_count = sum(1 for _ in snake_path(int(start_x), int(start_y), int(end_x), int(end_y)))
    elapsed = []
    total_jpeg_bytes = 0

    capture = VideoCapture(cam)
    capture.set_format(CAMERA_FRAME_WIDTH, CAMERA_FRAME_HEIGHT, CAMERA_FRAME_PIXEL_FORMAT)
    with capture:
        capture_iterator = iter(capture)
        for i, (x, y) in enumerate(snake_path(int(start_x), int(start_y), int(end_x), int(end_y))):
            start_t = time.time()
            stage.move_absolute(x, y)
            for _ in range(CAMERA_FRAME_SKIP):
                next(capture_iterator)
            frame = next(capture_iterator)
            with (output_directory / f"{i:05d}_x{stage.x}_y{stage.y}.jpg").open("wb") as out:
                out.write(frame.data)
                total_jpeg_bytes += len(frame.data)
            end_t = time.time()
            elapsed.append(end_t - start_t)
            if len(elapsed) > 16:
                elapsed.pop(0)
            avg_per_it = fmean(elapsed)
            print(
                f"\r{int((i + 1)/path_count*100):03d}%, "
                f"{i + 1}/{path_count} captures, "
                f"{avg_per_it:.2f} s/it, "
                f"ETA {timedelta(seconds=int((path_count-i)*avg_per_it))}, "
                f"{total_jpeg_bytes/(1024*1024):.2f} MiB JPEG",
                end="")

    cam.close()

if __name__ == "__main__":
    main()