# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

"""
This module contains classes to wrap Python VTK to make nice molecular plots.
"""


__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__date__ = "Nov 27, 2011"


import os
import itertools
import math
import subprocess
import time

import numpy as np
try:
    import vtk
    from vtk import vtkInteractorStyleTrackballCamera
except ImportError:
    # VTK not present. The Camera is to set object to avoid errors in unittest.
    vtk = None
    vtkInteractorStyleTrackballCamera = object

from monty.serialization import loadfn
from monty.dev import requires

from pymatgen.util.coord_utils import in_coord_list
from pymatgen.core.periodic_table import Specie
from pymatgen.core.structure import Structure
from pymatgen.core.sites import PeriodicSite


module_dir = os.path.dirname(os.path.abspath(__file__))
EL_COLORS = loadfn(os.path.join(module_dir, "ElementColorSchemes.yaml"))


class StructureVis(object):
    """
    Provides Structure object visualization using VTK.
    """

    @requires(vtk, "Visualization requires the installation of VTK with "
                   "Python bindings.")
    def __init__(self, element_color_mapping=None, show_unit_cell=True,
                 show_bonds=False, show_polyhedron=True,
                 poly_radii_tol_factor=0.5, excluded_bonding_elements=None):
        """
        Constructs a Structure Visualization.

        Args:
            element_color_mapping: Optional color mapping for the elements,
                as a dict of {symbol: rgb tuple}. For example, {"Fe": (255,
                123,0), ....} If None is specified, a default based on
                Jmol"s color scheme is used.
            show_unit_cell: Set to False to not show the unit cell
                boundaries. Defaults to True.
            show_bonds: Set to True to show bonds. Defaults to True.
            show_polyhedron: Set to True to show polyhedrons. Defaults to
                False.
            poly_radii_tol_factor: The polyhedron and bonding code uses the
                ionic radii of the elements or species to determine if two
                atoms are bonded. This specifies a tolerance scaling factor
                such that atoms which are (1 + poly_radii_tol_factor) * sum
                of ionic radii apart are still considered as bonded.
            excluded_bonding_elements: List of atom types to exclude from
                bonding determination. Defaults to an empty list. Useful
                when trying to visualize a certain atom type in the
                framework (e.g., Li in a Li-ion battery cathode material).

        Useful keyboard shortcuts implemented.
            h : Show help
            A/a : Increase/decrease cell by one unit vector in a-direction
            B/b : Increase/decrease cell by one unit vector in b-direction
            C/c : Increase/decrease cell by one unit vector in c-direction
            # : Toggle showing of polyhedrons
            - : Toggle showing of bonds
            [ : Decrease poly_radii_tol_factor by 0.05
            ] : Increase poly_radii_tol_factor by 0.05
            r : Reset camera direction
            o : Orthogonalize structure
            Up/Down : Rotate view along Up direction by 90 clock/anticlockwise
            Left/right : Rotate view along camera direction by 90
            clock/anticlockwise
        """
        # create a rendering window and renderer
        self.ren = vtk.vtkRenderer()
        self.ren_win = vtk.vtkRenderWindow()
        self.ren_win.AddRenderer(self.ren)
        self.ren.SetBackground(1, 1, 1)
        self.title = "Structure Visualizer"
        # create a renderwindowinteractor
        self.iren = vtk.vtkRenderWindowInteractor()
        self.iren.SetRenderWindow(self.ren_win)
        self.mapper_map = {}
        self.structure = None

        if element_color_mapping:
            self.el_color_mapping = element_color_mapping
        else:
            self.el_color_mapping = EL_COLORS["VESTA"]
        self.show_unit_cell = show_unit_cell
        self.show_bonds = show_bonds
        self.show_polyhedron = show_polyhedron
        self.poly_radii_tol_factor = poly_radii_tol_factor
        self.excluded_bonding_elements = excluded_bonding_elements if \
            excluded_bonding_elements else []
        self.show_help = True
        self.supercell = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        self.redraw()

        style = StructureInteractorStyle(self)
        self.iren.SetInteractorStyle(style)
        self.ren.parent = self

    def rotate_view(self, axis_ind=0, angle=0):
        """
        Rotate the camera view.

        Args:
            axis_ind: Index of axis to rotate. Defaults to 0, i.e., a-axis.
            angle: Angle to rotate by. Defaults to 0.
        """
        camera = self.ren.GetActiveCamera()
        if axis_ind == 0:
            camera.Roll(angle)
        elif axis_ind == 1:
            camera.Azimuth(angle)
        else:
            camera.Pitch(angle)
        self.ren_win.Render()

    def write_image(self, filename="image.png", magnification=1,
                    image_format="png"):
        """
        Save render window to an image.

        Arguments:
            filename:
                filename to save to. Defaults to image.png.
            magnification:
                magnification. Use it to render high res images.
            image_format:
                choose between jpeg, png.  Png is the default.
        """
        render_large = vtk.vtkRenderLargeImage()
        render_large.SetInput(self.ren)
        if image_format == "jpeg":
            writer = vtk.vtkJPEGWriter()
            writer.SetQuality(80)
        else:
            writer = vtk.vtkPNGWriter()

        render_large.SetMagnification(magnification)
        writer.SetFileName(filename)

        writer.SetInputConnection(render_large.GetOutputPort())
        self.ren_win.Render()
        writer.Write()
        del render_large

    def redraw(self, reset_camera=False):
        """
        Redraw the render window.

        Args:
            reset_camera: Set to True to reset the camera to a
                pre-determined default for each structure.  Defaults to False.
        """
        self.ren.RemoveAllViewProps()
        self.picker = None
        self.add_picker_fixed()
        self.helptxt_mapper = vtk.vtkTextMapper()
        tprops = self.helptxt_mapper.GetTextProperty()
        tprops.SetFontSize(14)
        tprops.SetFontFamilyToTimes()
        tprops.SetColor(0, 0, 0)

        if self.structure is not None:
            self.set_structure(self.structure, reset_camera)

        self.ren_win.Render()

    def orthongonalize_structure(self):
        if self.structure is not None:
            self.set_structure(self.structure.copy(sanitize=True))
        self.ren_win.Render()

    def display_help(self):
        """
        Display the help for various keyboard shortcuts.
        """
        helptxt = ["h : Toggle help",
                   "A/a, B/b or C/c : Increase/decrease cell by one a,"
                   " b or c unit vector", "# : Toggle showing of polyhedrons",
                   "-: Toggle showing of bonds", "r : Reset camera direction",
                   "[/]: Decrease or increase poly_radii_tol_factor "
                   "by 0.05. Value = " + str(self.poly_radii_tol_factor),
                   "Up/Down: Rotate view along Up direction by 90 "
                   "clockwise/anticlockwise",
                   "Left/right: Rotate view along camera direction by "
                   "90 clockwise/anticlockwise", "s: Save view to image.png",
                   "o: Orthogonalize structure"]
        self.helptxt_mapper.SetInput("\n".join(helptxt))
        self.helptxt_actor.SetPosition(10, 10)
        self.helptxt_actor.VisibilityOn()

    def set_structure(self, structure, reset_camera=True, to_unit_cell=True):
        """
        Add a structure to the visualizer.

        Args:
            structure: structure to visualize
            reset_camera: Set to True to reset the camera to a default
                determined based on the structure.
            to_unit_cell: Whether or not to fall back sites into the unit cell.
        """
        self.ren.RemoveAllViewProps()

        has_lattice = hasattr(structure, "lattice")

        if has_lattice:
            s = Structure.from_sites(structure, to_unit_cell=to_unit_cell)
            s.make_supercell(self.supercell, to_unit_cell=to_unit_cell)
        else:
            s = structure

        inc_coords = []
        for site in s:
            self.add_site(site)
            inc_coords.append(site.coords)

        count = 0
        labels = ["a", "b", "c"]
        colors = [(1, 0, 0), (0, 1, 0), (0, 0, 1)]

        if has_lattice:
            matrix = s.lattice.matrix

        if self.show_unit_cell and has_lattice:
            #matrix = s.lattice.matrix
            self.add_text([0, 0, 0], "o")
            for vec in matrix:
                self.add_line((0, 0, 0), vec, colors[count])
                self.add_text(vec, labels[count], colors[count])
                count += 1
            for (vec1, vec2) in itertools.permutations(matrix, 2):
                self.add_line(vec1, vec1 + vec2)
            for (vec1, vec2, vec3) in itertools.permutations(matrix, 3):
                self.add_line(vec1 + vec2, vec1 + vec2 + vec3)

        if self.show_bonds or self.show_polyhedron:
            elements = sorted(s.composition.elements, key=lambda a: a.X)
            anion = elements[-1]

            def contains_anion(site):
                for sp in site.species_and_occu.keys():
                    if sp.symbol == anion.symbol:
                        return True
                return False

            anion_radius = anion.average_ionic_radius
            for site in s:
                exclude = False
                max_radius = 0
                color = np.array([0, 0, 0])
                for sp, occu in site.species_and_occu.items():
                    if sp.symbol in self.excluded_bonding_elements \
                            or sp == anion:
                        exclude = True
                        break
                    max_radius = max(max_radius, sp.average_ionic_radius)
                    color = color + \
                            occu * np.array(self.el_color_mapping.get(sp.symbol,
                                                                      [0, 0, 0]))

                if not exclude:
                    max_radius = (1 + self.poly_radii_tol_factor) * \
                                 (max_radius + anion_radius)
                    nn = structure.get_neighbors(site, float(max_radius))
                    nn_sites = []
                    for nnsite, dist in nn:
                        if contains_anion(nnsite):
                            nn_sites.append(nnsite)
                            if not in_coord_list(inc_coords, nnsite.coords):
                                self.add_site(nnsite)
                    if self.show_bonds:
                        self.add_bonds(nn_sites, site)
                    if self.show_polyhedron:
                        color = [i / 255 for i in color]
                        self.add_polyhedron(nn_sites, site, color)

        if self.show_help:
            self.helptxt_actor = vtk.vtkActor2D()
            self.helptxt_actor.VisibilityOn()
            self.helptxt_actor.SetMapper(self.helptxt_mapper)
            self.ren.AddActor(self.helptxt_actor)
            self.display_help()

        camera = self.ren.GetActiveCamera()
        if reset_camera:
            if has_lattice:
                #Adjust the camera for best viewing
                lengths = s.lattice.abc
                pos = (matrix[1] + matrix[2]) * 0.5 + \
                      matrix[0] * max(lengths) / lengths[0] * 3.5
                camera.SetPosition(pos)
                camera.SetViewUp(matrix[2])
                camera.SetFocalPoint((matrix[0] + matrix[1] + matrix[2]) * 0.5)
            else:
                origin = s.center_of_mass
                max_site = max(
                    s, key=lambda site: site.distance_from_point(origin))
                camera.SetPosition(origin + 5 * (max_site.coords - origin))
                camera.SetFocalPoint(s.center_of_mass)

        self.structure = structure
        self.title = s.composition.formula

    def zoom(self, factor):
        """
        Zoom the camera view by a factor.
        """
        camera = self.ren.GetActiveCamera()
        camera.Zoom(factor)
        self.ren_win.Render()

    def show(self):
        """
        Display the visualizer.
        """
        self.iren.Initialize()
        self.ren_win.SetSize(800, 800)
        self.ren_win.SetWindowName(self.title)
        self.ren_win.Render()
        self.iren.Start()

    def add_site(self, site):
        """
        Add a site to the render window. The site is displayed as a sphere, the
        color of which is determined based on the element. Partially occupied
        sites are displayed as a single element color, though the site info
        still shows the partial occupancy.

        Args:
            site: Site to add.
        """
        start_angle = 0
        radius = 0
        total_occu = 0

        for specie, occu in site.species_and_occu.items():
            radius += occu * (specie.ionic_radius
                              if isinstance(specie, Specie)
                                 and specie.ionic_radius
                              else specie.average_ionic_radius)
            total_occu += occu

        vis_radius = 0.2 + 0.002 * radius

        for specie, occu in site.species_and_occu.items():
            if not specie:
                color = (1, 1, 1)
            elif specie.symbol in self.el_color_mapping:
                color = [i / 255 for i in self.el_color_mapping[specie.symbol]]
            mapper = self.add_partial_sphere(site.coords, vis_radius, color,
                                             start_angle, start_angle + 360 * occu)
            self.mapper_map[mapper] = [site]
            start_angle += 360 * occu

        if total_occu < 1:
            mapper = self.add_partial_sphere(site.coords, vis_radius, (1,1,1),
                                             start_angle, start_angle + 360 * (1 - total_occu))
            self.mapper_map[mapper] = [site]

    def add_partial_sphere(self, coords, radius, color, start=0, end=360,
                           opacity=1.0):
        sphere = vtk.vtkSphereSource()
        sphere.SetCenter(coords)
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(18)
        sphere.SetPhiResolution(18)
        sphere.SetStartTheta(start)
        sphere.SetEndTheta(end)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(sphere.GetOutputPort())
        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        actor.GetProperty().SetColor(color)
        actor.GetProperty().SetOpacity(opacity)
        self.ren.AddActor(actor)
        return mapper

    def add_text(self, coords, text, color=(0, 0, 0)):
        """
        Add text at a coordinate.

        Args:
            coords: Coordinates to add text at.
            text: Text to place.
            color: Color for text as RGB. Defaults to black.
        """
        source = vtk.vtkVectorText()
        source.SetText(text)
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(source.GetOutputPort())
        follower = vtk.vtkFollower()
        follower.SetMapper(mapper)
        follower.GetProperty().SetColor(color)
        follower.SetPosition(coords)
        follower.SetScale(0.5)
        self.ren.AddActor(follower)
        follower.SetCamera(self.ren.GetActiveCamera())

    def add_line(self, start, end, color=(0.5, 0.5, 0.5), width=1):
        """
        Adds a line.

        Args:
            start: Starting coordinates for line.
            end: Ending coordinates for line.
            color: Color for text as RGB. Defaults to grey.
            width: Width of line. Defaults to 1.
        """
        source = vtk.vtkLineSource()
        source.SetPoint1(start)
        source.SetPoint2(end)

        vertexIDs = vtk.vtkStringArray()
        vertexIDs.SetNumberOfComponents(1)
        vertexIDs.SetName("VertexIDs")
        # Set the vertex labels
        vertexIDs.InsertNextValue("a")
        vertexIDs.InsertNextValue("b")
        source.GetOutput().GetPointData().AddArray(vertexIDs)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(source.GetOutputPort())
        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        actor.GetProperty().SetColor(color)
        actor.GetProperty().SetLineWidth(width)
        self.ren.AddActor(actor)

    def add_polyhedron(self, neighbors, center, color, opacity=1.0,
                       draw_edges=False, edges_color=[0.0, 0.0, 0.0],
                       edges_linewidth=2):
        """
        Adds a polyhedron.

        Args:
            neighbors: Neighbors of the polyhedron (the vertices).
            center: The atom in the center of the polyhedron.
            color: Color for text as RGB.
            opacity: Opacity of the polyhedron
            draw_edges: If set to True, the a line will be drawn at each edge
            edges_color: Color of the line for the edges
            edges_linewidth: Width of the line drawn for the edges
        """
        points = vtk.vtkPoints()