# Python script to generate Patterson map animations for the presentation
# at the CCP4 Study Weekend 2001.
#
# Software required to run this script (this software is freely available):
#     Python 1.5.2, http://www.python.org/
#     ReportLab 1.00, http://www.reportlab.com/
#
# conversion of resulting PDF file to GIF (for import in PowerPoint):
#   Photoshop -> File -> Open *.pdf, Resolution = 150 pixels/inch
#                File -> Save for Web

import math
import copy
from reportlab.pdfgen import canvas
from reportlab.lib.units import inch

golden_ratio = (math.sqrt(5.) + 1.) / 2.

mol_gray = 0.8
shift_mol_gray = 0.6
pat_gray = 0.3
uc_gray = 0.5

class UnitCellCanvas(canvas.Canvas):

  def __init__(self, filename):
    self.px = 3.4 * inch
    self.py = (3 * golden_ratio + 0.4) * inch
    canvas.Canvas.__init__(self, filename, pagesize = (self.py, self.px))
    self.closing = 0
    self.setup_coordiante_system()

  def setup_coordiante_system(self):
    self.translate(0.2 * inch, self.px - 0.2 * inch)
    self.rotate(-90)
    self.scale(1. * inch, golden_ratio * inch)

  def showPage(self):
    if (not self.closing):
      canvas.Canvas.showPage(self)
      self.setup_coordiante_system()

  def save(self):
    self.closing = 1
    canvas.Canvas.save(self)

  def real_circle(self, x, y, radius, stroke = 1, fill = 1):
    self.saveState()
    self.translate(x, y)
    self.scale(1, 1. / golden_ratio)
    self.circle(0, 0, radius, stroke, fill)
    self.restoreState()


def patterson(atoms):
  vects = []
  for a1 in atoms:
    for a2 in atoms:
      vects.append((a2[0] - a1[0], a2[1] - a1[1]))
  return vects

def mark_inter_vects(c, vects, color, radius, alternate = 0, eps = 0.001):
  c.setFillColorRGB(color[0], color[1], color[2])
  for ux in xrange(-2, 4):
    for uy in xrange(-2, 4):
      if (alternate):
        if ((ux + uy) % 2 == 0):
          c.setFillColorRGB(1, 0, 0)
        else:
          c.setFillColorRGB(1, 0, 1)
      for v in vects:
        x = ux + v[0]
        y = uy + v[1]
        if (0. - eps <= x <= 3. + eps and 0. - eps <= y <= 3. + eps):
          c.real_circle(x, y, radius, stroke = 0, fill = 1)

def mark_intra_vects(c, vects, color, radius):
  c.setFillColorRGB(color[0], color[1], color[2])
  for v in vects:
    c.real_circle(1 + v[0], 1 + v[1], radius, stroke = 0, fill = 1)

def light_color(color):
  light = [0] * 3
  for i in xrange(3): light[i] = max(0.4, color[i])
  return light

def draw_patterson(c, structure, apply_lattice_translations = 0, gray = 0,
                   radius = 0.024, alternate = 0):
  for mol in structure:
    if (gray): color = (pat_gray, pat_gray, pat_gray)
    else:      color = mol[0]
    atoms = mol[1]
    vects = patterson(atoms)
    if (apply_lattice_translations):
      if (gray): mark_inter_vects(c, vects, color, radius)
      else: mark_inter_vects(c, vects, light_color(color), radius, alternate)
    mark_intra_vects(c, vects, color, radius)

def draw_unit_cells(c, color = (0.2, 0.5, 0.3), linewidth = 0.02):
  c.setStrokeColorRGB(color[0], color[1], color[2])
  c.setLineWidth(linewidth)
  c.grid(range(4), range(4))

def draw_molecule(c, uc, mol, gray = 0):
  color = mol[0]
  atoms = mol[1]
  if (gray):
    c.setStrokeColorRGB(shift_mol_gray, shift_mol_gray, shift_mol_gray)
  else:
    c.setStrokeColorRGB(color[0], color[1], color[2])
  c.setFillColorRGB(mol_gray, mol_gray, mol_gray)
  p = c.beginPath()
  d = p.moveTo
  for atom in atoms:
    d(uc[0] + atom[0], uc[1] + atom[1])
    d = p.lineTo
  p.close()
  c.drawPath(p, stroke = 1, fill = not gray)
  if (gray): c.setFillColorRGB(1.0, 0.0, 0.0)
  else:      c.setFillColorRGB(0.0, 0.0, 0.0)
  for atom in atoms:
    c.real_circle(uc[0] + atom[0], uc[1] + atom[1], 0.02, stroke = 0, fill = 1)

def draw_structure(c, structure):
  c.setLineWidth(0.01)
  for mol in structure:
    for ux in xrange(3):
      for uy in xrange(3):
        draw_molecule(c , (ux, uy), mol)

def merge_molecules(structure):
  merge_color = [0] * 3
  merge_atoms = []
  for mol in structure:
    color = mol[0]
    atoms = mol[1]
    for i in xrange(3):
      merge_color[i] = min(merge_color[i] + color[i], 1.)
    for atom in atoms:
      merge_atoms.append(atom)
  merge_color = [0, 1, 0]
  return [[merge_color, merge_atoms]]

def scale_structure(structure, scale_factor):
  ss = copy.deepcopy(structure)
  for mol in ss:
    for atom in mol[1]:
      for i in xrange(2): atom[i] = atom[i] * scale_factor
  return ss

def shift_structure(structure, offs):
  ss = copy.deepcopy(structure)
  for mol in ss:
    for atom in mol[1]:
      for i in xrange(2): atom[i] = atom[i] - offs[i]
  return ss

def draw_shrunken_and_obs_super_imposed(c, merge_structure, scale_factor, ss):
  # shrunken first intra-vectors + apply_lattice_translations
  # with observed patterson super-imposed
  draw_unit_cells(c)
  draw_patterson(c, ss, apply_lattice_translations = 1, alternate = 1)
  c.translate(1, 1)
  c.scale(scale_factor, scale_factor)
  c.translate(-1, -1)
  draw_unit_cells(c, color = (uc_gray, uc_gray, uc_gray))
  draw_patterson(c, merge_structure, apply_lattice_translations = 1, gray = 1)

if (__name__ == '__main__'):
  structure = [
    [
      (1, 0, 0),
      [
        [6, 10],
        [16, 4],
        [22, 16],
        [7, 19],
      ]
    ],
    [
      (0, 0, 1),
      [
        [3, 27],
        [19, 21],
        [12, 37],
      ]
    ],
  ]
  for mol in structure:
    for atom in mol[1]:
      atom[0] = atom[0] / 25.
      atom[1] = atom[1] / 40.

  merge_structure = merge_molecules(structure)

  c = UnitCellCanvas('plot.pdf')

  # grid
  draw_unit_cells(c)
  c.showPage()

  # grid + first molecule
  draw_unit_cells(c)
  draw_structure(c, structure[0:1])
  c.showPage()

  # grid + first molecule + shifted vectors for each atom in first molecule
  for atom1 in structure[0][1]:
    for atom2 in structure[0][1]:
      ss = shift_structure(structure, atom2)
      draw_unit_cells(c)
      draw_structure(c, structure[0:1])
      draw_molecule(c, (1, 1), ss[0], gray = 1)
      if (atom1 == atom2): break
    c.showPage()

  # grid + first molecule + intra-vectors
  draw_unit_cells(c)
  draw_structure(c, structure[0:1])
  draw_patterson(c, structure[0:1])
  c.showPage()

  # grid + first molecule + intra-vectors + second molecule
  draw_unit_cells(c)
  draw_structure(c, structure)
  draw_patterson(c, structure[0:1])
  c.showPage()

  # grid + first molecule + intra-vectors + second molecule + intra-vectors
  draw_unit_cells(c)
  draw_structure(c, structure)
  draw_patterson(c, structure)
  c.showPage()

  # grid + first molecule + intra-vectors + second molecule + intra-vectors
  # + cross-vectors
  draw_unit_cells(c)
  draw_structure(c, structure)
  draw_patterson(c, merge_structure)
  draw_patterson(c, structure)
  c.showPage()

  # grid + first intra-vectors + second intra-vectors + cross-vectors
  draw_unit_cells(c)
  draw_patterson(c, merge_structure)
  draw_patterson(c, structure)
  c.showPage()

  # grid + first intra-vectors + second intra-vectors + cross-vectors
  # + apply_lattice_translations
  draw_unit_cells(c)
  draw_patterson(c, merge_structure, apply_lattice_translations = 1)
  draw_patterson(c, structure, apply_lattice_translations = 1)
  c.showPage()

  # grid + first intra-vectors + second intra-vectors + cross-vectors
  # + apply_lattice_translations all in gray
  draw_unit_cells(c)
  draw_patterson(c, merge_structure, apply_lattice_translations = 1, gray = 1)
  c.showPage()

  # grid + first molecule + intra-vectors
  draw_unit_cells(c)
  draw_structure(c, structure[0:1])
  draw_patterson(c, structure[0:1])
  c.showPage()

  # shrink first molecule + intra-vectors + apply_lattice_translations
  for scale_factor in (1.0, 0.81, 0.62, 0.43):
    ss = scale_structure(structure[0:1], scale_factor)
    draw_unit_cells(c)
    draw_structure(c, ss)
    draw_patterson(c, ss, apply_lattice_translations = 1, alternate = 1)
    c.showPage()

  scale_factor = 0.43
  ss = scale_structure(structure[0:1], scale_factor)

  # shrunken first intra-vectors + apply_lattice_translations
  draw_unit_cells(c)
  draw_patterson(c, ss, apply_lattice_translations = 1, alternate = 1)
  c.showPage()

  integration_radius = 0.33
  integration_omit_radius = 0.12

  # shrunken first intra-vectors + apply_lattice_translations
  # show integration radius
  c.setFillColorRGB(1, 1, 0)
  c.real_circle(1, 1, integration_radius, stroke = 0, fill = 1)
  draw_unit_cells(c)
  draw_patterson(c, ss, apply_lattice_translations = 1, alternate = 1)
  c.showPage()

  # observed map super-imposed
  c.setFillColorRGB(1, 1, 0)
  c.real_circle(1, 1, integration_radius, stroke = 0, fill = 1)
  draw_shrunken_and_obs_super_imposed(c, merge_structure, scale_factor, ss)
  c.showPage()

  # show integration shell
  c.setFillColorRGB(1, 1, 0)
  c.real_circle(1, 1, integration_radius, stroke = 0, fill = 1)
  c.setFillColorRGB(1, 1, 1)
  c.real_circle(1, 1, 0.12, stroke = 0, fill = 1)
  draw_shrunken_and_obs_super_imposed(c, merge_structure, scale_factor, ss)
  c.showPage()

  # grid + first molecule + intra-vectors
  draw_unit_cells(c)
  draw_structure(c, structure[0:1])
  draw_patterson(c, structure[0:1])
  c.showPage()

  # grid + first intra-vectors + first inter-vectors
  draw_unit_cells(c)
  draw_patterson(c, structure[0:1], apply_lattice_translations = 1)
  c.showPage()

  # grid + first molecule + intra-vectors + obvserved patterson super-imposed
  draw_unit_cells(c)
  draw_patterson(c, structure[0:1], apply_lattice_translations = 1,
                 radius = 0.03)
  draw_patterson(c, merge_structure, apply_lattice_translations = 1, gray = 1)
  c.showPage()

  c.save()