# -*- coding: utf-8 -*-
"""
Visualize Atoms and Cavities with GR3
"""

from math import sin, cos, sqrt, pi
import gr3
from ..config.configuration import config
from ..core import calculation as calculation
import numpy as np
import numpy.linalg as la
import os
from ctypes import c_int
from ..util.gl_util import (
    create_perspective_projection_matrix,
    create_look_at_matrix,
    create_rotation_matrix_homogenous,
    create_translation_matrix_homogenous,
)


class Visualization(object):
    """
    Visualize Atoms and Cavities with GR3
    """

    def __init__(self):
        self.max_side_lengths = 1.0
        self.mat = np.eye(4)
        self.d = self.max_side_lengths * 2
        self.pos = np.array((0, 0, self.d))
        self.up = np.array((0, 1, 0))
        self.right = np.array((1, 0, 0))

        self.near = 0.1
        self.far = 3 * self.d
        self.fov = 45.0

        self.width = 0
        self.height = 0
        self.usecurrentframebuffer = False
        #  self.assigned_opengl_context = None

        self.results = None
        self.settings = VisualizationSettings()
        self.objectids = [None]

    def setresults(self, results):
        """
        Set the results that will be displayed

        **Parameters:**
            `results` :
                :class:`core.data.Results` object to visualize
        """
        self.results = results
        max_side_lengths = max(results.atoms.volume.side_lengths)
        self.d = self.d / self.max_side_lengths * max_side_lengths
        self.max_side_lengths = max_side_lengths
        self.far = 6 * max_side_lengths
        self.create_scene()
        if self.width != 0 and self.height != 0:
            self.paint(self.width, self.height)
        self.settings.visible_domain_indices = None
        self.settings.visible_surface_cavity_indices = None
        self.settings.visible_center_cavity_indices = None

    def create_scene(self):
        """
        Create GR3 meshes. ``self.results`` contains the mesh data
        and in ``self.settings`` is specified which meshes shound be rendered.
        """

        c = config.Colors.background
        gr3.setbackgroundcolor(c[0], c[1], c[2], 1.0)
        gr3.clear()

        if self.results is None:
            return

        show_domains = self.settings.show_domains
        show_surface_cavities = self.settings.show_surface_cavities
        show_center_cavities = self.settings.show_center_cavities
        if show_center_cavities and self.results.center_cavities is not None:
            show_surface_cavities = False
        elif show_surface_cavities and self.results.surface_cavities is not None:
            show_domains = False

        self.objectids = [None]
        edges = self.results.atoms.volume.edges
        num_edges = len(edges)
        edge_positions = [edge[0] for edge in edges]
        edge_directions = [
            [edge[1][i] - edge[0][i] for i in range(3)] for edge in edges
        ]
        edge_lengths = [sum([c * c for c in edge]) ** 0.5 for edge in edge_directions]
        edge_radius = min(edge_lengths) / 200
        if self.settings.show_bounding_box:
            gr3.drawcylindermesh(
                num_edges,
                edge_positions,
                edge_directions,
                [config.Colors.bounding_box] * num_edges,
                [edge_radius] * num_edges,
                edge_lengths,
            )
            corners = list(
                set(
                    [tuple(edge[0]) for edge in edges]
                    + [tuple(edge[1]) for edge in edges]
                )
            )
            num_corners = len(corners)
            gr3.drawspheremesh(
                num_corners,
                corners,
                [config.Colors.bounding_box] * num_corners,
                [edge_radius] * num_corners,
            )

        if self.settings.show_atoms and self.results.atoms is not None:
            visible_atom_indices = self.settings.visible_atom_indices
            if visible_atom_indices is not None:
                visible_atom_indices = [
                    comp
                    for comp in visible_atom_indices
                    if 0 <= comp < self.results.atoms.number
                ]
            else:
                visible_atom_indices = range(self.results.atoms.number)
            if len(visible_atom_indices) == 0:
                visible_atom_indices = None
            if visible_atom_indices is not None:
                visible_atom_indices = np.array(visible_atom_indices)
                gr3.drawspheremesh(
                    len(visible_atom_indices),
                    self.results.atoms.positions[visible_atom_indices],
                    self.results.atoms.colors[visible_atom_indices],
                    np.ones(len(visible_atom_indices)) * config.OpenGL.atom_radius,
                )
                if self.settings.show_bonds:
                    bonds = self.results.atoms.bonds
                    for start_index, target_indices in enumerate(bonds):
                        if start_index not in visible_atom_indices:
                            continue
                        target_indices = np.array(
                            [i for i in target_indices if i in visible_atom_indices]
                        )
                        if len(target_indices) == 0:
                            continue
                        start_position = self.results.atoms.positions[start_index]
                        target_positions = self.results.atoms.positions[target_indices]
                        directions = target_positions - start_position
                        bond_lengths = la.norm(directions, axis=1)
                        directions /= bond_lengths.reshape(len(directions), 1)
                        gr3.drawcylindermesh(
                            len(target_indices),
                            target_positions,
                            -directions,
                            [config.Colors.bonds] * self.results.atoms.number,
                            np.ones(bond_lengths.shape) * config.OpenGL.bond_radius,
                            bond_lengths,
                        )

        if self.results is None:
            return
        if show_domains and self.results.domains is not None:
            self.draw_cavities(
                self.results.domains,
                config.Colors.domain,
                "domain",
                self.settings.visible_domain_indices,
            )
        if show_surface_cavities and self.results.surface_cavities is not None:
            self.draw_cavities(
                self.results.surface_cavities,
                config.Colors.surface_cavity,
                "surface cavity",
                self.settings.visible_surface_cavity_indices,
            )
        if show_center_cavities and self.results.center_cavities is not None:
            self.draw_cavities(
                self.results.center_cavities,
                config.Colors.center_cavity,
                "center cavity",
                self.settings.visible_center_cavity_indices,
            )

    def draw_cavities(self, cavities, color, cavity_type, indices=None):
        if indices is None:
            indices = range(self.results.domains.number)
        for index in set(indices):
            if 0 <= index < len(cavities.triangles):
                triangles = cavities.triangles[index]
                gr3._gr3.gr3_setobjectid(gr3.c_int(len(self.objectids)))
                self.objectids.append((cavity_type, index))
                mesh = gr3.createmesh(
                    triangles.shape[1] * 3,
                    triangles[0, :, :, :],
                    triangles[1, :, :, :],
                    [color] * (triangles.shape[1] * 3),
                )
                gr3.drawmesh(
                    mesh, 1, (0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 1, 1), (1, 1, 1)
                )
                gr3.deletemesh(c_int(mesh.value))
        gr3._gr3.gr3_setobjectid(gr3.c_int(0))

    def zoom(self, delta):
        """
        Adjust the zoom level.

        **Parameters:**
            `delta` :
                the increment of the distance from the camera to the center
        """
        zoom_v = 1.0 / 20
        zoom_cap = self.d + zoom_v * delta < (self.max_side_lengths) * 4
        if self.d + zoom_v * delta > 0 and zoom_cap:
            self.d += zoom_v * delta

    def translate_mouse(self, dx, dy):
        """
        Translate the model.
        """
        trans_v = 1.0 / 1000
        trans_factor = trans_v * max(self.d, 20)
        m = create_translation_matrix_homogenous(
            -dx * trans_factor, dy * trans_factor, 0
        )
        self.mat = self.mat.dot(m)

    def set_focus_on(self, x, y, z):
        self.mat[:3, 3] = (x, y, z)

    def rotate_mouse(self, dx, dy):
        """
        Rotate the model according to a mouse movement (dx, dy) on the screen.
        """
        rot_v = 1.0 / 13000
        if dx == 0 and dy == 0:
            return
        rot_axis = np.array((-dy, -dx, 0.0))
        rot_factor = max(self.d, 20) * rot_v * la.norm(rot_axis)
        # rotation matrix with min rotation angle
        m = create_rotation_matrix_homogenous(
            rot_factor, rot_axis[0], rot_axis[1], rot_axis[2]
        )
        self.mat = self.mat.dot(m)

    def reset_view(self):
        self.mat = np.eye(4)

    def set_camera(self, width, height):
        """
        Update the shown scene after the perspective has changed.
        """
        rightt = self.mat[:3, 0]
        upt = self.mat[:3, 1]
        pt = self.mat[:3, 2] * self.d
        t = self.mat[:3, 3]

        self.proj_mat = create_perspective_projection_matrix(
            np.radians(self.fov), 1.0 * width / height, self.near, self.far
        )
        gr3.setcameraprojectionparameters(self.fov, self.near, self.far)
        self.lookat_mat = create_look_at_matrix(pt + t, t, upt)
        gr3.cameralookat(
            pt[0] + t[0],
            pt[1] + t[1],
            pt[2] + t[2],
            t[0],
            t[1],
            t[2],
            upt[0],
            upt[1],
            upt[2],
        )

    def assign_opengl_context(self, context):
        self.assigned_opengl_context = context

    def paint(self, width, height, usecurrentframebuffer=None, device_pixel_ratio=1):
        """
        Refresh the OpenGL scene.
        """
        self.width = width
        self.height = height
        self.set_camera(width, height)
        if usecurrentframebuffer is not None:
            self.usecurrentframebuffer = usecurrentframebuffer
        if self.usecurrentframebuffer:
            gr3.usecurrentframebuffer()
        gr3.drawimage(
            0,
            width * device_pixel_ratio,
            0,
            height * device_pixel_ratio,
            width,
            height,
            gr3.GR3_Drawable.GR3_DRAWABLE_OPENGL,
        )

    def save_screenshot(
        self, file_name, width=3840, height=2160, first=True, last=True
    ):
        """
        Save a screenshot in the given resolution.
        ``first`` and ``last`` can be used to indicate if the first or last screenshot is taken when
        multiple images are saved in a loop for example.
        """
        gr3.export(file_name, width, height)

    def get_object_at_2dposition(self, x, y):
        oid = gr3.c_int(0)
        gr3._gr3.gr3_selectid(
            gr3.c_int(x),
            gr3.c_int(y),
            gr3.c_int(self.width),
            gr3.c_int(self.height),
            gr3.byref(oid),
        )
        return self.objectids[oid.value]

    def get_object_at_3dposition(self, x, y, z):
        if self.results is None:
            return None
        if self.results.atoms is not None:
            positions = self.results.atoms.positions
            distances = la.norm(positions - (x, y, z), axis=1)
            nearest_atom_index = np.argmin(distances)

            # calculate atom radius
            edges = self.results.atoms.volume.edges
            edge_directions = [
                [edge[1][i] - edge[0][i] for i in range(3)] for edge in edges
            ]
            edge_lengths = [
                sum([c * c for c in edge]) ** 0.5 for edge in edge_directions
            ]
            edge_radius = min(edge_lengths) / 200
            atom_radius = 4 * edge_radius

            if 0.95 < distances[nearest_atom_index] / atom_radius < 1.05:
                return ("atom", nearest_atom_index)


class VisualizationSettings(object):
    """
    Container for settings that control which part of the data is visualized.
    Attributes:

        `show_domains`

        `show_surface_cavities`

        `show_center_cavities`

        `show_atoms`

        `show_bounding_box`

    """

    def __init__(
        self,
        domains=False,
        show_surface_cavities=True,
        show_center_cavities=False,
        atoms=True,
        bonds=True,
        bounding_box=True,
        visible_atom_indices=None,
        visible_domain_indices=None,
        visible_surface_cavity_indices=None,
        visible_center_cavity_indices=None,
    ):
        self.show_domains = domains
        self.visible_domain_indices = visible_domain_indices
        self.show_surface_cavities = show_surface_cavities
        self.visible_surface_cavity_indices = visible_surface_cavity_indices
        self.show_center_cavities = show_center_cavities
        self.visible_center_cavity_indices = visible_center_cavity_indices
        self.show_atoms = atoms
        self.visible_atom_indices = visible_atom_indices
        self.show_bonds = bonds
        self.show_bounding_box = bounding_box
