#
#  Project: MXCuBE
#  https://github.com/mxcube
#
#  This file is part of MXCuBE software.
#
#  MXCuBE is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  MXCuBE is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU General Lesser Public License
#  along with MXCuBE. If not, see <http://www.gnu.org/licenses/>.

__copyright__ = """2019 by the MXCuBE collaboration """
__license__ = "LGPLv3+"

import base64
import copy
from functools import reduce
from io import BytesIO

import numpy as np
from PIL import Image

from mxcubecore import HardwareRepository as HWR
from mxcubecore.HardwareObjects.abstract.AbstractSampleView import (
    AbstractSampleView,
    ShapeState,
)
from mxcubecore.model import queue_model_objects


def combine_images(img1, img2):
    if img1.size != img2.size:
        raise ValueError("Images must be the same size")

    combined_img = Image.new("RGB", img1.size)

    pixels1 = img1.load()
    pixels2 = img2.load()
    combined_pixels = combined_img.load()

    width, height = img1.size
    for x in range(width):
        for y in range(height):
            pixel1 = pixels1[x, y]
            pixel2 = pixels2[x, y]

            if pixel2[0] <= 200 and pixel2[1] <= 60 and pixel2[2] <= 140:
                combined_pixels[x, y] = pixel1
            else:
                combined_pixels[x, y] = pixel2

    return combined_img


class SampleView(AbstractSampleView):
    def __init__(self, name):
        AbstractSampleView.__init__(self, name)
        self._shapes = {}

    def init(self):
        super(SampleView, self).init()
        self._last_oav_image = None

        self.hide_grid_threshold = self.get_property("hide_grid_threshold", 5)
        for motor_name, motor_ho in HWR.beamline.diffractometer.get_motors().items():
            if motor_ho:
                motor_ho.connect("stateChanged", self._update_shape_positions)

    def _update_shape_positions(self, *args, **kwargs):
        for shape in self.get_shapes():
            shape.update_position(HWR.beamline.diffractometer.motor_positions_to_screen)

        self.emit("shapesChanged")

    @property
    def shapes(self):
        return self._shapes

    def start_centring(self, tree_click=True):
        """
        Starts centring procedure
        """
        pass

    def cancel_centring(self):
        """
        Cancels current centring procedure
        """
        pass

    def start_auto_centring(self):
        """
        Start automatic centring procedure
        """
        pass

    def get_snapshot(self, overlay=None, bw=False, return_as_array=False):
        """
        Get snapshot(s)

        Args:
            overlay(str): Image data with shapes and other items to display on the snapshot
            bw(bool): return grayscale image
            return_as_array(bool): return as np array

        Returns:
            (BytesIO) snapshot as bytes image
        """
        img = self.take_snapshot(overlay_data=overlay, bw=bw)

        if return_as_array:
            return np.array(img)

        buffered = BytesIO()
        img.save(buffered, format="JPEG")

        return buffered

    def save_snapshot(self, path, overlay=None, bw=False):
        """
        Save a snapshot to file.

        Args:
            path (str): The filename.
            overlay(str): Image data with shapes and other items to display on the snapshot
            bw(bool): return grayscale image
        """
        img = self.take_snapshot(overlay_data=overlay, bw=bw)
        img.save(path)

        self._last_oav_image = path

    def take_snapshot(self, overlay_data=None, bw=False):
        """
        Get snapshot with overlaid data.

        Args:
            overlay_data (str): base64 encoded image to lay over camera image
            bw (bool): return grayscale image

        Returns:
            (Image) rgb or grayscale image
        """
        data, width, height = self.camera.get_last_image()

        img = Image.frombytes("RGB", (width, height), data)

        if overlay_data:
            overlay_data = base64.b64decode(overlay_data)
            overlay_image = Image.open(BytesIO(overlay_data))
            overlay_image = overlay_image.resize(
                (width, height), Image.Resampling.LANCZOS
            )
            img = combine_images(img, overlay_image.convert("RGB"))

        if bw:
            img.convert("1")

        return img

    def get_last_image_path(self):
        return self._last_oav_image

    def add_shape(self, shape):
        """
        Add the shape <shape> to the dictionary of handled shapes.

        Args:
            param (shape): Shape to add.
            type (shape): Shape object.
        """
        self.shapes[shape.id] = shape
        shape.shapes_hw_object = self

    def add_shape_from_mpos(
        self,
        mpos_list,
        screen_coord,
        t,
        state: ShapeState = "SAVED",
        user_state: ShapeState = "SAVED",
    ):
        """
        Adds a shape of type <t>, with motor positions from mpos_list and
        screen position screen_coord.

        Args:
            mpos_list (list[mpos_list]): List of motor positions
            screen_coord (tuple(x, y): Screen coordinate for shape
            t (str): Type str for shape, P (Point), L (Line), G (Grid)

        Returns:
            (Shape) Shape of type <t>
        """
        cls_dict = {"P": Point, "L": Line, "G": Grid, "2DP": TwoDPoint}
        _cls = cls_dict[t]
        shape = None

        if _cls:
            shape = _cls(mpos_list, screen_coord)
            # In case the shape is being recreated, we need to restore it's state.
            shape.state = state
            shape.user_state = user_state

            self.add_shape(shape)

        return shape

    def add_shape_from_refs(
        self, refs, t, state: ShapeState = "SAVED", user_state: ShapeState = "SAVED"
    ):
        """
        Adds a shape of type <t>, taking motor positions and screen positions
        from reference points in refs.

        Args:
            refs (list[str]): List of id's of the reference Points
            t (str): Type str for shape, P (Point), L (Line), G (Grid)

        Returns:
            (Shape): Shape of type <t>
        """
        mpos = [self.get_shape(refid).mpos() for refid in refs]
        spos_list = [self.get_shape(refid).screen_coord for refid in refs]
        spos = reduce((lambda x, y: tuple(x) + tuple(y)), spos_list, ())
        shape = self.add_shape_from_mpos(mpos, spos, t, state, user_state)
        shape.refs = refs

        return shape

    def delete_shape(self, sid):
        """
        Removes the shape with id <sid> from the list of handled shapes.

        Args:
            sid (str): The id of the shape to remove

        Returns:
            (Shape): The removed shape
        """
        shape = self.shapes.pop(sid, None)

        if shape:
            shape.shapes_hw_object = None

        return shape

    def select_shape(self, sid):
        """
        Select the shape <shape>.

        Args:
            sid (str): Id of the shape to select.
        """
        shape = self.shapes.get(sid, None)

        if shape:
            shape.select()

    def de_select_shape(self, sid):
        """
        De-select the shape with id <sid>.

        Args:
            sid (str): The id of the shape to de-select.
        """
        shape = self.shapes.get(sid, None)

        if shape:
            shape.de_select()

    def is_selected(self, sid):
        """
        Check if Shape with <sid> is selected.

        Returns:
            (Boolean) True if Shape with <sid> is selected False otherwise
        """
        shape = self.shapes.get(sid, None)
        return bool(shape and shape.is_selected())

    def get_selected_shapes(self):
        """
        Get all selected shapes.

        Returns:
           (list[Shape]) List of selected Shapes
        """
        return [s for s in self.shapes.values() if s.is_selected()]

    def de_select_all(self):
        """De select all shapes."""

        for shape in self.shapes.values():
            shape.de_select()

    def select_shape_with_cpos(self, cpos):
        """
        Selects shape with the assocaitaed centered position <cpos>

        Args:
            cpos (CenteredPosition)
        """
        return

    def clear_all(self):
        """
        Clear the shapes, remove all contents.
        """
        self._shapes = {}
        Grid.SHAPE_COUNT = 0
        Line.SHAPE_COUNT = 0
        Point.SHAPE_COUNT = 0

    def get_shapes(self):
        """
        Get all Shapes.

        Returns:
            (list[Shape]) All the shapes
        """
        return self.shapes.values()

    def get_points(self):
        """
        Get all Points currently handled.

        Returns:
            (list[Point]) All points currently handled
        """
        current_points = []

        for shape in self.get_shapes():
            if isinstance(shape, Point):
                current_points.append(shape)

        return current_points

    def get_lines(self):
        """
        Get all Lines currently handled.

        Returns:
            (list[Line]) All lines currently handled
        """
        lines = []

        for shape in self.get_shapes():
            if isinstance(shape, Line):
                lines.append(shape)

        return lines

    def get_grids(self):
        """
        Get all Grids currently handled.

        Returns:
            (list[Grid]) All lines currently handled
        """
        grid = []

        for shape in self.get_shapes():
            if isinstance(shape, Grid):
                grid.append(shape)

        return grid

    def get_shape(self, sid):
        """
        Get Shape with id <sid>.

        Args:
            sid (str): id of Shape to retrieve

        Returns:
            (Shape) All the shapes
        """
        return self.shapes.get(sid, None)

    # For backwards compatibility with old ShapeHisotry object
    # returns first of selected grids
    def get_grid(self):
        """
        Get the first of the selected grids, (the one that was selected first in
        a sequence of select operations)

        Returns:
            (dict): The first selected grid as a dictionary
        """
        grid = None

        for shape in self.get_shapes():
            if isinstance(shape, Grid):
                grid = shape.as_dict()
                break

        return grid

    def set_grid_data(self, sid, result_data, data_file_path):
        """
        Sets grid rsult data for a shape with the specified id.

        Args:
            sid (str): The id of the shape to set grid data for.
            result_data: The result data to set for the shape. Either a base64 encoded string for PNG/image
            or a dictionary for RGB (keys are cell number and value RGBa list). Data is only updated if result is RGB based
            data_file_path (str): The path to the data file associated with the result data.

        Returns:
            None

        Raises:
            AttributeError: If no shape with the specified id exists.
        """

        shape = self.get_shape(sid)

        if shape:
            if shape.result and type(shape.result) == dict:
                # append data
                shape.result.update(result_data)
            else:
                shape.set_result(result_data)
                shape.result_data_path = data_file_path

            self.emit("newGridResult", shape)
        else:
            msg = "Cant set result for %s, no shape with id %s" % (sid, sid)
            raise AttributeError(msg)

    def get_grid_data(self, key):
        result = {}
        shape = self.get_shape(key)

        if shape:
            result = shape.get_result()

        return result

    def inc_used_for_collection(self, cpos):
        """
        Increase counter that keepts on collect made on this shape,
        shape with associated CenteredPosition cpos

        Args:
            cpos (CenteredPosition): CenteredPosition of shape
        """
        pass


class Shape(object):
    """
    Base class for shapes.
    """

    SHAPE_COUNT = 0

    def __init__(self, mpos_list=[], screen_coord=(-1, -1)):
        object.__init__(self)
        Shape.SHAPE_COUNT += 1
        self.t = "S"
        self.id = ""
        self.cp_list = []
        self.name = ""
        self.state: ShapeState = "SAVED"
        self.user_state: ShapeState = "SAVED"  # used to persist user preferences in regards whether to show or hide particular shape.
        self.label = ""
        self.screen_coord = screen_coord
        self.selected = False
        self.refs = []
        self.shapes_hw_object = None

        self.add_cp_from_mp(mpos_list)

    def get_centred_positions(self):
        """
        :returns: The centred position(s) associated with the shape.
        :rtype: List of CentredPosition objects.
        """
        return self.cp_list

    def get_centred_position(self):
        return self.get_centred_positions()[0]

    def select(self):
        self.selected = True

    def de_select(self):
        self.selected = False

    def is_selected(self):
        return self.selected

    def update_position(self, transform):
        spos_list = [transform(cp.as_dict()) for cp in self.cp_list]
        spos_list = tuple([pos for l in spos_list for pos in l])
        self.screen_coord = spos_list

    def add_cp_from_mp(self, mpos_list):
        for mp in mpos_list:
            self.cp_list.append(queue_model_objects.CentredPosition(mp))

    def set_id(self, id_num):
        self.id = self.t + "%s" % id_num
        self.name = self.label + "-%s" % id_num

    def move_to_mpos(self, mpos_list, screen_coord=[]):
        self.cp_list = []
        self.add_cp_from_mp(mpos_list)

        if screen_coord:
            self.screen_coord = screen_coord

    def update_from_dict(self, shape_dict):
        # We don't allow id or result updates
        shape_dict.pop("id", None)
        shape_dict.pop("result", None)

        for key, value in shape_dict.items():
            if hasattr(self, key):
                setattr(self, key, value)

    def as_dict(self):
        cpos_list = []

        for cpos in self.cp_list:
            cpos_list.append(cpos.as_dict())

        d = copy.deepcopy(vars(self))

        # Do not serialize Shapes HW Object
        d.pop("shapes_hw_object")

        # replace cpos_list with a list of motor positions
        d.pop("cp_list")
        d["motor_positions"] = str(cpos_list)

        return d


class Point(Shape):
    SHAPE_COUNT = 0

    def __init__(self, mpos_list, screen_coord):
        Shape.__init__(self, mpos_list, screen_coord)
        Point.SHAPE_COUNT += 1
        self.t = "P"
        self.label = "Point"
        self.set_id(Point.SHAPE_COUNT)

    def mpos(self):
        return self.cp_list[0].as_dict()

    def set_id(self, id_num):
        Shape.set_id(self, id_num)
        self.cp_list[0].index = self.name

    def as_dict(self):
        d = Shape.as_dict(self)
        # replace cpos_list with the motor positions
        d["motor_positions"] = self.cp_list[0].as_dict()
        return d


class TwoDPoint(Point):
    SHAPE_COUNT = 0

    def __init__(self, mpos_list, screen_coord):
        Point.__init__(self, mpos_list, screen_coord)
        self.t = "2DP"
        self.label = "2D-Point"
        self.set_id(Point.SHAPE_COUNT)


class Line(Shape):
    SHAPE_COUNT = 0

    def __init__(self, mpos_list, screen_coord):
        Shape.__init__(self, mpos_list, screen_coord)
        Line.SHAPE_COUNT += 1
        self.t = "L"
        self.label = "Line"
        self.set_id(Line.SHAPE_COUNT)

    def set_id(self, id_num):
        Shape.set_id(self, id_num)
        self.cp_list[0].index = self.name

    def get_centred_positions(self):
        return [self.start_cpos, self.end_cpos]

    def get_points_index(self):
        if all(self.cp_list):
            return (self.cp_list[0].get_index(), self.cp_list[1].get_index())


class Grid(Shape):
    SHAPE_COUNT = 0

    def __init__(self, mpos_list, screen_coord):
        Shape.__init__(self, mpos_list, screen_coord)
        Grid.SHAPE_COUNT += 1
        self.t = "G"
        self.set_id(Grid.SHAPE_COUNT)

        self.width = -1
        self.height = -1
        self.cell_count_fun = "zig-zag"
        self.cell_h_space = -1
        self.cell_height = -1
        self.cell_v_space = -1
        self.cell_width = -1
        self.label = "Grid"
        self.num_cols = -1
        self.num_rows = -1
        self.selected = False
        # result is a base64 encoded string for PNG/image heatmap results
        # or a dictionary (for RGB number based results)
        self.result = None
        self.pixels_per_mm = [1, 1]
        self.beam_pos = [1, 1]
        self.beam_width = 0
        self.beam_height = 0
        self.hide_threshold = 5

        self.set_id(Grid.SHAPE_COUNT)

    def update_position(self, transform):
        phi_pos = HWR.beamline.diffractometer.omega.get_value() % 360
        _d = abs((self.get_centred_position().phi % 360) - phi_pos)

        if self.user_state == "HIDDEN":
            self.state = "HIDDEN"
            return

        if min(_d, 360 - _d) > self.shapes_hw_object.hide_grid_threshold:
            self.state = "HIDDEN"
        else:
            super(Grid, self).update_position(transform)
            self.state = "SAVED"

    def get_centred_position(self):
        return self.cp_list[0]

    def get_grid_range(self):
        return (
            float(self.cell_width * (self.num_cols - 1)),
            float(self.cell_height * (self.num_rows - 1)),
        )

    def get_num_lines(self):
        if self.cell_count_fun == "zig-zag":
            return self.num_rows
        elif self.cell_count_fun == "inverse-zig-zag":
            return self.num_cols
        else:
            return self.num_rows

    def set_id(self, id_num):
        Shape.set_id(self, id_num)
        self.cp_list[0].index = self.name

    def set_result(self, result_data):
        self.result = result_data

    def get_result(self):
        return self.result

    def as_dict(self):
        d = Shape.as_dict(self)
        # replace cpos_list with the motor positions
        d["motor_positions"] = self.cp_list[0].as_dict()

        pixels_per_mm = HWR.beamline.diffractometer.get_pixels_per_mm()
        beam_pos = HWR.beamline.beam.get_beam_position_on_screen()
        size_x, size_y, shape, _label = HWR.beamline.beam.get_value()

        d["x1"] = -float((beam_pos[0] - d["screen_coord"][0]) / pixels_per_mm[0])
        d["y1"] = -float((beam_pos[1] - d["screen_coord"][1]) / pixels_per_mm[1])
        d["steps_x"] = d["num_cols"]
        d["steps_y"] = d["num_rows"]
        d["dx_mm"] = d["width"] / pixels_per_mm[0]
        d["dy_mm"] = d["height"] / pixels_per_mm[1]
        d["beam_width"] = size_x
        d["beam_height"] = size_y
        d["angle"] = 0

        return d
