import argparse
import os
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
import skimage


from textual import on
from textual.app import App, ComposeResult
from textual.containers import Horizontal, HorizontalScroll, Vertical, VerticalScroll
from textual.widgets import Button, Footer, Header, Label, Select, RadioButton, RadioSet, Switch

from _imagepanel import ImagePanel
from _slider import Slider

__version__: str = '2025.4.4'  # this must be here for pyproject.toml to find it.


class climax(App):
    """a Command Line IMAge eXplorer."""

    __version__: str = __version__  # defines climax.__version__
    __default_cmap__: str = 'gray'
    __default_channels__: tuple[str] = ('', )

    CSS_PATH = "climax.tcss"
    orientations = {'x-y': [0, 1, 2], 'z-y': [2, 1, 0], 'z-x': [1, 2, 0]}
    image_extensions: tuple[str] = ('.tif', '.tiff', '.jpg', '.jpeg', '.gif', '.png', '.bmp')
    zoom_factors: tuple[float] = (1.0, 2.0, 4.0, 0.25, 0.50)
    colormaps: list[str] = plt.colormaps()


    BINDINGS = [("k", "toggle_dark", "toggle dark mode"),
                (".", "next_slice", "next slice"),
                ("comma", "prev_slice", "previous slice"),
                ("z", "zoom", "zoom"),
                ]

    AUTO_FOCUS = "#image-panel"

    def __init__(self, filename: str|list[str]='tests/cells_movie.tif', channel_strs: Optional[tuple[str]] = ('', ), cmap: str = 'gray', figsize: tuple[int, int] =(4,4)):
        super().__init__()
        
        # This if statement here for debugging purposes.
        if filename==[] or filename is None or filename==False:
            filename = 'tests/cells_movie.tif'

        if isinstance(filename, str):
            volume = climax.read_image(filename, channel_strs)
            if volume is None:
                # throw exception instead.
                print(f'No volume found at that path. The current path is {os.path.curdir}.')

        elif isinstance(filename, list):
            volume_list: list = []
            for afile in filename:
                loaded_stack = climax.read_image(afile, channel_strs)
                if loaded_stack is not None:
                    volume_list.append(loaded_stack)

            volume = np.concatenate(np.asarray(volume_list), axis=volume_list[0].ndim-1)

        self.volume: np.ndarray = volume
        self.slices: np.ndarray = None

        match self.volume.ndim:
            case 1:
                self.volume = np.expand_dims(self.volume, axis=(0, 1, 2, 3, ))
            case 2:
                self.volume = np.expand_dims(self.volume, axis=(0, 1, 2, ))
            case 3:
                self.volume = np.expand_dims(self.volume, axis=(0, 1, ))
            case 4:
                self.volume = np.expand_dims(self.volume, axis=0)
            case 5:
                pass
            case _:
                print(f"I don't know how to deal with {self.volume.ndim}-dimensional images. But file an issue at https://bitbucket.org/rfg_lab/climax/issues and someone will teach me.")

       # image quality control: for images that are not unsigned int, convert to unsigned int with as many bits as necessary for the image range.
        if not self.volume.dtype.kind == 'u':
            self.volume = self.volume.astype(np.uint16)
            im_range = int(np.ceil(self.volume.max()-self.volume.min())) if self.volume.min() < 0. else int(np.ceil(self.volume.max()))
            min_bytes = int(np.ceil(len(bin(im_range)[2:])/8))
            self.volume = self.volume.astype(np.dtype(f'uint{min_bytes*8}'))

        self.color = self.volume[0]
        self.slices = self.color[0]

        self.cmap: str = cmap 
        self.figsize: tuple[int, int] = figsize  # not implemented
        self.pix_val_min_max: list = [np.min(self.color), np.max(self.color)]
        self.invert_image: bool = False  # invert image
        self.continuous_update: bool = True  # continuous update
        self.curslice: int = 0
        self.view_plane: str = list(climax.orientations.keys())[0]
        self.zoom_index: int = 0
        self.image_data = np.transpose(self.slices, climax.orientations[self.view_plane])

        self.create_UI()
        self._setupUI()
        self.image_panel = ImagePanel(self.image_data[self.curslice], cmap=self.cmap, id="image-panel")

    def create_UI(self):
        max_Z = self.slices.shape[0]-1 if self.slices.shape[0]>1 else 1        
        self.Z_slider_name_label = Label("image slice:", id="image-slice-label")
        self.Z_slider_value_label = Label(f"{self.curslice}", id="image-slice-value")
        self.Z_slider = Slider(0, max_Z, id="Z-slider")
        self.Z_slider_group = HorizontalScroll(self.Z_slider_name_label, self.Z_slider, self.Z_slider_value_label, id="Z-slider-group")

        max_time = self.volume.shape[1]-1 if self.volume.shape[1]>1 else 1
        self.time_slider_name_label = Label("time point: ", id="time-point-label")
        self.time_slider_value_label = Label(f"{self.curslice}", id="time-point-value")
        self.time_slider = Slider(0, max_time, id="time-slider")
        self.time_slider_group = HorizontalScroll(self.time_slider_name_label, self.time_slider, self.time_slider_value_label, id="time-slider-group")

        max_channel = self.volume.shape[0]-1 if self.volume.shape[0]>1 else 1
        self.channel_slider_name_label = Label("channel:    ", id="channel-point-label")
        self.channel_slider_value_label = Label(f"{self.curslice}", id="channel-point-value")
        self.channel_slider = Slider(0, max_channel, id="channel-slider")
        self.channel_slider_group = HorizontalScroll(self.channel_slider_name_label, self.channel_slider, self.channel_slider_value_label, id="channel-slider-group")

        self.vmin_label = Label("minimum pixel value:", id="vmin-label")
        self.vmin_value_label = Label(f"{self.curslice}", id="vmin-value-label")
        self.vmin_slider = Slider(0, self.pix_val_min_max[1], value=self.pix_val_min_max[0], id="vmin-slider")
        self.vmin_group = HorizontalScroll(self.vmin_label, self.vmin_slider, self.vmin_value_label, id="vmin-slider-group")

        self.vmax_label = Label("maximum pixel value:", id="vmax-label")
        self.vmax_value_label = Label(f"{self.curslice}", id="vmax-value-label")
        self.vmax_slider = Slider(0, self.pix_val_min_max[1], value=self.pix_val_min_max[1], id="vmax-slider")
        self.vmax_group = HorizontalScroll(self.vmax_label, self.vmax_slider, self.vmax_value_label, id="vmax-slider-group")

        self.rotate_button = Button("rotate", id="rotate-button", tooltip="rotate 90° clockwise")
        self.flip_horizontal_button = Button("flip h.", id="fliph-button", tooltip="flip horizontal")
        self.flip_vertical_button = Button("flip v.", id="flipv-button", tooltip="flip vertical")
        self.invert_image_label = Label("invert: ", id="invert-label")
        self.invert_image_switch = Switch(value=self.invert_image, id="invert-switch", tooltip="invert image")
        self.continuous_update_label = Label("update: ", id="update-label")
        self.continuous_update_switch = Switch(value=self.continuous_update, id="update-switch", tooltip="update continuously")

        self.slice_radio_set = RadioSet(*[RadioButton(anorientation, value=(not anindex)) for anindex, anorientation in enumerate(self.orientations.keys())], id="plane-radio", tooltip='slicing plane')
                
        self.auto_contrast_button = Button("auto", id="auto-contrast-button", tooltip="Scale display from image mode to 99th percentile")
        self.colormap_select = Select.from_values(self.colormaps, value=self.cmap, id="colormap-select")

        self.button_bar = HorizontalScroll(self.rotate_button, self.flip_horizontal_button, self.flip_vertical_button, self.invert_image_label, self.invert_image_switch, self.continuous_update_label, self.continuous_update_switch)

        self.dashboard = Horizontal(Vertical(self.channel_slider_group, self.time_slider_group, self.Z_slider_group, self.button_bar), self.slice_radio_set, Vertical(self.vmin_group, self.vmax_group, HorizontalScroll(self.auto_contrast_button, self.colormap_select)), id='dashboard')

    def _setupUI(self):
        self.channel_slider.disabled = False if self.volume.shape[0] > 1 else True  # IS THERE A WAY TO GRAY THESE OUT?
        self.time_slider.disabled = False if self.volume.shape[1] > 1 else True
        self.Z_slider.disabled = False if self.slices.shape[0] > 1 else True
        self.slice_radio_set.disabled = False if self.slices.shape[0] > 1 else True


    def display_slice(self, reset_slices: Optional[bool] = False, curslice: Optional[int] = 0) -> None:
        if reset_slices or self.view_plane != str(self.slice_radio_set.pressed_button.label):
            self.view_plane = str(self.slice_radio_set.pressed_button.label)
            
            self.image_data = np.transpose(self.slices, climax.orientations[self.view_plane])
    
            self.Z_slider.max = self.image_data.shape[0]-1

            self.curslice = curslice if 0 <= curslice <= self.Z_slider.max else 0
            self.Z_slider.value = self.curslice

            self._setupUI()

        elif curslice:
            self.curslice = curslice if 0 <= curslice <= self.Z_slider.max else self.curslice
            self.Z_slider.value = self.curslice


        self.image_panel.set_pixels(self.image_data[self.curslice], self.vmin_slider.value, self.vmax_slider.value, zoom_factor=self.zoom_factors[self.zoom_index])
        self.image_panel.display_slice()
            
    def compose(self) -> ComposeResult:
        """Create child widgets for the app."""
        yield Header(id="header")

        image_scroll = VerticalScroll(
            HorizontalScroll(self.image_panel, id="horizontal-scroll"),
            id="vertical-scroll"
        )
        yield Vertical(self.dashboard, image_scroll, id="main")
        yield Footer(id="footer")


    def action_toggle_dark(self) -> None:
        """An action to toggle dark mode."""
        self.theme = (
            "textual-dark" if self.theme == "textual-light" else "textual-light"
        )

    def action_next_slice(self):
        self.display_slice(curslice=self.curslice+1)

    def action_prev_slice(self):
        self.display_slice(curslice=self.curslice-1)

    def action_zoom(self) -> None:
        """An action to toggle dark mode."""
        self.zoom_index = (self.zoom_index + 1) % len(self.zoom_factors)
        self.display_slice()

    @on(Slider.Changed, "#Z-slider")
    def _Z_slider_change(self) -> None:
        if not self.Z_slider.disabled:
            self.curslice = self.Z_slider.value
            self.Z_slider_value_label.update(f"{self.curslice}")

            if self.continuous_update:
                self.display_slice()

    @on(Slider.Clicked, "#Z-slider")
    def _Z_slider_click(self) -> None:
        if not self.Z_slider.disabled and not self.continuous_update:
            self.display_slice()

    @on(Slider.Keyed, "#Z-slider")
    def _Z_slider_key(self) -> None:
        if not self.Z_slider.disabled:
            self.display_slice()

    @on(Slider.Changed, "#time-slider")
    def _time_slider_change(self) -> None:
        if not self.time_slider.disabled:
            curtime = self.time_slider.value
            self.slices = self.color[curtime]
            self.time_slider_value_label.update(f"{curtime}")
            if self.continuous_update:
                self.display_slice(True, self.curslice)

    @on(Slider.Keyed, "#time-slider")
    def _time_slider_key(self) -> None:
        if not self.time_slider.disabled:
            curtime = self.time_slider.value
            self.slices = self.color[curtime]
            self.time_slider_value_label.update(f"{curtime}")
            self.display_slice(True, self.curslice)

    @on(Slider.Clicked, "#time-slider")
    def _time_slider_click(self) -> None:
        if not self.time_slider.disabled and not self.continuous_update:
            self.display_slice(True, self.curslice)

    @on(Slider.Changed, "#channel-slider")
    def _channel_slider_change(self) -> None:
        if not self.channel_slider.disabled:
            curchannel = self.channel_slider.value
            self.color = self.volume[curchannel]
            self.slices = self.color[self.time_slider.value]
            self.channel_slider_value_label.update(f"{curchannel}")
            if self.continuous_update:
                self.display_slice(True, self.curslice)

    @on(Slider.Keyed, "#channel-slider")
    def _channel_slider_key(self) -> None:
        if not self.channel_slider.disabled:
            curchannel = self.channel_slider.value
            self.color = self.volume[curchannel]
            self.slices = self.color[self.time_slider.value]
            self.channel_slider_value_label.update(f"{curchannel}")
            self.display_slice(True, self.curslice)

    @on(Slider.Clicked, "#channel-slider")
    def _channel_slider_click(self) -> None:
        if not self.channel_slider.disabled and not self.continuous_update:
            self.display_slice(True, self.curslice)

    @on(Slider.Changed, "#vmin-slider")
    def _vmin_slider_change(self) -> None:
        if self.vmin_slider.value > self.vmax_slider.value:
            self.vmax_slider.value = self.vmin_slider.value
        self.vmin_value_label.update(f"{self.vmin_slider.value}")
        if self.continuous_update:
            self.display_slice()

    @on(Slider.Keyed, "#vmin-slider")
    def _vmin_slider_key(self) -> None:
        if self.vmin_slider.value > self.vmax_slider.value:
            self.vmax_slider.value = self.vmin_slider.value
        self.vmin_value_label.update(f"{self.vmin_slider.value}")
        self.display_slice()

    @on(Slider.Clicked, "#vmin-slider")
    def _vmin_slider_click(self) -> None:
            if not self.continuous_update:
                self.display_slice()

    @on(Slider.Changed, "#vmax-slider")
    def _vmax_slider_change(self) -> None:
        if self.vmax_slider.value < self.vmin_slider.value:
            self.vmin_slider.value = self.vmax_slider.value
        self.vmax_value_label.update(f"{self.vmax_slider.value}")
        if self.continuous_update:
            self.display_slice()

    @on(Slider.Keyed, "#vmax-slider")
    def _vmax_slider_key(self) -> None:
        if self.vmax_slider.value < self.vmin_slider.value:
            self.vmin_slider.value = self.vmax_slider.value
        self.vmax_value_label.update(f"{self.vmax_slider.value}")
        self.display_slice()

    @on(Slider.Clicked, "#vmax-slider")
    def _vmax_slider_click(self) -> None:
            if not self.continuous_update:
                self.display_slice()

    @on(Button.Pressed, "#rotate-button")
    def _rotate_button_press(self) -> None:
        self.volume = np.rot90(self.volume, -1, (3, 4))
        self.color = self.volume[self.channel_slider.value]
        self.slices = self.color[self.time_slider.value]
        self.display_slice(True, self.curslice)

    @on(Button.Pressed, "#fliph-button")
    def _flip_horizontal_button_press(self) -> None:
        self.volume = self.volume[..., ::-1]
        self.color = self.volume[self.channel_slider.value]
        self.slices = self.color[self.time_slider.value]
        self.display_slice(True, self.curslice)

    @on(Button.Pressed, "#flipv-button")
    def _flip_vertical_button_press(self) -> None:
        self.volume = self.volume[..., ::-1, :]
        self.color = self.volume[self.channel_slider.value]
        self.slices = self.color[self.time_slider.value]
        self.display_slice(True, self.curslice)

    @on(Switch.Changed, "#update-switch")
    def _update_switch_change(self) -> None:
        self.continuous_update = not self.continuous_update

    @on(Switch.Changed, "#invert-switch")
    def _invert_switch_change(self) -> None:
        self.invert_image = not self.invert_image
        if self.cmap.endswith("_r"):
            self.cmap = self.cmap[:-2]
        else:
            self.cmap += "_r"

        self.colormap_select.value = self.cmap
        self.image_panel.set_cmap(self.cmap)
        self.display_slice()

    @on(Button.Pressed, "#auto-contrast-button")
    def _auto_contrast_pressed(self) -> None:
        """Scale from mode to 99th percentile."""
        low = climax.mode(self.image_data[self.curslice,:,:].ravel())[0]
        high = int(np.percentile(self.image_data[self.curslice,:,:], 99))

        self.vmin_slider.value = low  # This triggers vmin and vmax changed events. Would be nicer if only one event was triggered, but we still need to update the sliders ...
        self.vmax_slider.value = high

        old_update = self.continuous_update
        self.continuous_update = True

        self._vmin_slider_change()
        self._vmax_slider_change()

        self.continuous_update = old_update

    @on(Select.Changed, "#colormap-select")
    def _cmap_changed(self) -> None:
        self.cmap = self.colormap_select.value
        self.image_panel.set_cmap(self.cmap)
        self.display_slice()

    @on(RadioSet.Changed, "#plane-radio")
    def _plane_changed(self) -> None:
        self.display_slice(True)

    @classmethod
    def read_image(cls, file_path: str, channel_strs: Optional[tuple[str]]=('',)) -> Optional[np.ndarray]:
        im: np.ndarray = None

        if os.path.isfile(file_path):        
            _, ext = os.path.splitext(file_path)

            # First try to read images within the allowed extensions.
            if str.lower(ext) in climax.image_extensions:
                im = skimage.io.imread(file_path)
            # If the extension is unknown, try to read as a tiff file.
            else:
                try:
                    im = skimage.io.imread(file_path, plugin='tifffile')
                except:
                    return None
        
        elif os.path.isdir(file_path):
            channels: list[np.ndarray] = [[] for _ in channel_strs]

            for filename in os.listdir(file_path):
                img = cls.read_image(os.path.join(file_path, filename))

                if img is not None:
                    index_list = [theindex for theindex in range(len(channel_strs)) if channel_strs[theindex] in filename]

                    if index_list != []:
                        channels[index_list[0]].append(img)

            im = np.asarray(channels)

        return im
    
    @classmethod
    def mode(cls, anarray: np.ndarray):
        vals, cnts = np.unique(anarray, return_counts=True)
        modes, counts = vals[cnts.argmax()], cnts.max()
        return modes[()], counts[()]
    
    @classmethod
    def parse_arguments(cls) -> argparse.Namespace:
        parser = argparse.ArgumentParser(
            prog='climax',
            description='a Comand Line IMAge eXplorer',
        )

        parser.add_argument('filename')  # this could be a list?
        parser.add_argument('-s', '--strings', nargs='*', default=climax.__default_channels__)
        parser.add_argument('-c', '--cmap', default=climax.__default_cmap__)

        return parser.parse_args()



if __name__ == "__main__":

    args = climax.parse_arguments()
    app = climax(args.filename, cmap=args.cmap, channel_strs=args.strings)
    app.run()