"""
Module for povray file format support.

See http://www.povray.org/ for details on the format.
"""
import os

import numpy as np

from ase.io.eps import EPS
from ase.data import chemical_symbols


def pa(array):
    """Povray array syntax"""
    return '<% 6.2f, % 6.2f, % 6.2f>' % tuple(array)


def pc(array):
    """Povray color syntax"""
    if type(array) == str:
        return 'color ' + array
    if type(array) == float:
        return 'rgb <%.2f>*3' % array
    if len(array) == 3:
        return 'rgb <%.2f, %.2f, %.2f>' % tuple(array)
    if len(array) == 4: # filter
        return 'rgbf <%.2f, %.2f, %.2f, %.2f>' % tuple(array)
    if len(array) == 5: # filter and transmit
        return 'rgbft <%.2f, %.2f, %.2f, %.2f, %.2f>' % tuple(array)


def get_bondpairs(atoms, radius=1.1):
    """Get all pairs of bonding atoms

    Return all pairs of atoms which are closer than radius times the
    sum of their respective covalent radii.  The pairs are returned as
    tuples::

      (a, b, (i1, i2, i3))

    so that atoms a bonds to atom b displaced by the vector::

        _     _     _
      i c + i c + i c ,
       1 1   2 2   3 3

    where c1, c2 and c3 are the unit cell vectors and i1, i2, i3 are
    integers."""
    
    from ase.data import covalent_radii
    from ase.calculators.neighborlist import NeighborList
    cutoffs = radius * covalent_radii[atoms.numbers]
    nl = NeighborList(cutoffs=cutoffs, self_interaction=False)
    nl.update(atoms)
    bondpairs = []
    for a in range(len(atoms)):
        indices, offsets = nl.get_neighbors(a)
        bondpairs.extend([(a, a2, offset)
                          for a2, offset in zip(indices, offsets)])
    return bondpairs


class POVRAY(EPS):
    default_settings = {
        # x, y is the image plane, z is *out* of the screen
        'display'      : True, # Display while rendering
        'pause'        : True, # Pause when done rendering (only if display)
        'transparent'  : True, # Transparent background
        'canvas_width' : None, # Width of canvas in pixels
        'canvas_height': None, # Height of canvas in pixels 
        'camera_dist'  : 50.,  # Distance from camera to front atom
        'image_plane'  : None, # Distance from front atom to image plane
        'camera_type'  : 'orthographic', # perspective, ultra_wide_angle
        'point_lights' : [],             # [[loc1, color1], [loc2, color2],...]
        'area_light'   : [(2., 3., 40.), # location
                          'White',       # color
                          .7, .7, 3, 3], # width, height, Nlamps_x, Nlamps_y
        'background'   : 'White',        # color
        'textures'     : None, # Length of atoms list of texture names
        'celllinewidth': 0.05, # Radius of the cylinders representing the cell
        'bondlinewidth': 0.10, # Radius of the cylinders representing the bonds
        'bondatoms'    : [], # [[atom1, atom2], ...] pairs of bonding atoms
        }

    def __init__(self, atoms, scale=1.0, **parameters):
        for k, v in self.default_settings.items():
            setattr(self, k, parameters.pop(k, v))
        EPS.__init__(self, atoms, scale=scale, **parameters)

    def cell_to_lines(self, A):
        return np.empty((0, 3)), None, None

    def write(self, filename, **settings):
        # Determine canvas width and height
        ratio = float(self.w) / self.h
        if self.canvas_width is None:
            if self.canvas_height is None:
                self.canvas_width = min(self.w * 15, 640)
            else:
                self.canvas_width = self.canvas_height * ratio
        elif self.canvas_height is not None:
            raise RuntimeError, "Can't set *both* width and height!"

        # Distance to image plane from camera
        if self.image_plane is None:
            if self.camera_type == 'orthographic':
                self.image_plane = 1 - self.camera_dist
            else:
                self.image_plane = 0
        self.image_plane += self.camera_dist

        # Produce the .ini file
        if filename.endswith('.pov'):
            ini = open(filename[:-4] + '.ini', 'w').write
        else:
            ini = open(filename + '.ini', 'w').write
        ini('Input_File_Name=%s\n' % filename)
        ini('Output_to_File=True\n')
        ini('Output_File_Type=N\n')
        ini('Output_Alpha=%s\n' % self.transparent)
        ini('; if you adjust Height, and width, you must preserve the ratio\n')
        ini('; Width / Height = %s\n' % repr(ratio))
        ini('Width=%s\n' % self.canvas_width)
        ini('Height=%s\n' % (self.canvas_width / ratio))
        ini('Antialias=True\n')
        ini('Antialias_Threshold=0.1\n')
        ini('Display=%s\n' % self.display)
        ini('Pause_When_Done=%s\n' % self.pause)
        ini('Verbose=False\n')
        del ini

        # Produce the .pov file
        w = open(filename, 'w').write
        w('#include "colors.inc"\n')
        w('#include "finish.inc"\n')
        w('\n')
        w('global_settings {assumed_gamma 1 max_trace_level 6}\n')
        w('background {%s}\n' % pc(self.background))
        w('camera {%s\n' % self.camera_type)
        w('  right -%.2f*x up %.2f*y\n' % (self.w, self.h))
        w('  direction %.2f*z\n' % self.image_plane)
        w('  location <0,0,%.2f> look_at <0,0,0>}\n' % self.camera_dist)
        for loc, rgb in self.point_lights:
            w('light_source {%s %s}\n' % (pa(loc), pc(rgb)))

        if self.area_light is not None:
            loc, color, width, height, nx, ny = self.area_light
            w('light_source {%s %s\n' % (pa(loc), pc(color)))
            w('  area_light <%.2f, 0, 0>, <0, %.2f, 0>, %i, %i\n' % (
                width, height, nx, ny))
            w('  adaptive 1 jitter}\n')

        w('\n')
        w('#declare simple = finish {phong 0.7}\n')
        w('#declare vmd = finish {'
          'ambient .0 '
          'diffuse .65 '
          'phong 0.1 '
          'phong_size 40. '
          'specular 0.500 }\n')
        w('#declare jmol = finish {'
          'ambient .2 '
          'diffuse .6 '
          'specular 1 '
          'roughness .001 '
          'metallic}\n')
        w('#declare ase2 = finish {'
          'ambient 0.05 '
          'brilliance 3 '
          'diffuse 0.6 '
          'metallic '
          'specular 0.70 '
          'roughness 0.04 '
          'reflection 0.15}\n')
        w('#declare ase3 = finish {'
          'ambient .15 '
          'brilliance 2 '
          'diffuse .6 '
          'metallic '
          'specular 1. '
          'roughness .001 '
          'reflection .0}\n')
        w('#declare glass = finish {' # Use filter 0.7
          'ambient .05 '
          'diffuse .3 '
          'specular 1. '
          'roughness .001}\n')
        w('#declare Rcell = %.3f;\n' % self.celllinewidth)
        w('#declare Rbond = %.3f;\n' % self.bondlinewidth)
        w('\n')
        w('#macro atom(LOC, R, COL, FIN)\n')
        w('  sphere{LOC, R texture{pigment{COL} finish{FIN}}}\n')
        w('#end\n')
        w('\n')
        
        z0 = self.X[:, 2].max()
        self.X -= (self.w / 2, self.h / 2, z0)

        # Draw unit cell
        if self.C is not None:
            self.C -= (self.w / 2, self.h / 2, z0)
            self.C.shape = (2, 2, 2, 3)
            for c in range(3):
                for j in ([0, 0], [1, 0], [1, 1], [0, 1]):
                    w('cylinder {')
                    for i in range(2):
                        j.insert(c, i)
                        w(pa(self.C[tuple(j)]) + ', ')
                        del j[c]
                    w('Rcell pigment {Black}}\n')

        # Draw atoms
        a = 0
        for loc, dia, color in zip(self.X, self.d, self.colors):
            tex = 'ase3'
            if self.textures is not None:
                tex = self.textures[a]
            w('atom(%s, %.2f, %s, %s) // #%i \n' % (
                pa(loc), dia / 2., pc(color), tex, a))
            a += 1

        # Draw atom bonds
        for pair in self.bondatoms:
            if len(pair) == 2:
                a, b = pair
                offset = (0, 0, 0)
            else:
                a, b, offset = pair
            R = np.dot(offset, self.A)
            mida = 0.5 * (self.X[a] + self.X[b] + R)
            midb = 0.5 * (self.X[a] + self.X[b] - R)
            if self.textures is not None:
                texa = self.textures[a]
                texb = self.textures[b]
            else:
                texa = texb = 'ase3'
            w('cylinder {%s, %s, Rbond texture{pigment {%s} finish{%s}}}\n' % (
                pa(self.X[a]), pa(mida), pc(self.colors[a]), texa))
            w('cylinder {%s, %s, Rbond texture{pigment {%s} finish{%s}}}\n' % (
                pa(self.X[b]), pa(midb), pc(self.colors[b]), texb))


def write_pov(filename, atoms, run_povray=False, **parameters):
    if isinstance(atoms, list):
        assert len(atoms) == 1
        atoms = atoms[0]
    assert 'scale' not in parameters
    POVRAY(atoms, **parameters).write(filename)
    if run_povray:
        errcode = os.system('povray %s.ini 2> /dev/null' % filename[:-4])
        if errcode != 0:
            raise OSError('Povray failed with error code %d' % errcode)