def usage():
  return """Python script to generate a synthetic photograph.
   Requirements:  Phenix must be installed and on path.
                  Mosflm must be installed and "ipmosflm" on path.

   Usage:
   1. First index the dataset using labelit.index.  Typical syntax to index two images is
      labelit.index /home/data/project/dataset1_1_E1_###.img 1 90

   2. Then

      labelit.precession_photo <command line arguments>

      Example:

      labelit.precession_photo bravais_choice=1 \\
      pixel_width=600 \\
      resolution_outer=3.0 \\
      intensity_full_scale=1000 \\
      plot_section="H,K,0" \\
      image_range=1,90 \\
      pdf_output.file="HKL.pdf"

   3. Command line arguments explained:

      labelit.precession_photo --help

   Known limitations:
   a) Since the source images are rotation photographs, there is an inherent
      approximation in assuming that the data was obtained at the central
      rotation angle (0.5 degree in a 0.0 to 1.0 rotation).
   b) Real precession photographs are obtained with a finite-width annulus;
      however the approximation here is one of an infinitesmal-width opening.
   c) This is a naive implementation that reads all of the data into memory
      first before calculating the image.  More efficient algorithms will be
      implemented later.

   Author: Nick Sauter, LBNL; May 8, 2010.
   Released under the CCTBX license; see cctbx.sf.net.
"""
import math,pickle,os
import libtbx.phil

precession_master_phil = libtbx.phil.parse("""
labelit_precession {
  pixel_width = 1500
    .type = int
    .help = "Width of pixel grid on which to calculate the synthesized image; must be sufficiently large to sample the Bragg spots"
  resolution_outer = 4.0
    .type = float
    .help = "The high-resolution limit (Angstroms) for the synthesized image"
  bravais_choice = None
    .type = int
    .help = "Crystal setting (integer ID#), as enumerated in labelit.index output, chosen for the synthesized section; print a list with labelit.stats_index"
  plot_section = None
    .type = str
    .help = "Miller indices of the section to be synthesized, comma separated, no spaces, such as H,0,L"
  image_range = None
    .type = ints(size=2)
    .help = "Range of data frames (inclusive) to be used for synthesized image, i.e., min frame no.,max frame no., no spaces"
  intensity_full_scale = 255
    .type = int
    .help = "Image intensity to be used as full scale (black)"
  pdf_output {
    file = None
      .type = str
      .help = "File name for pdf output"
  }
}
""")


from scitbx import matrix
from scitbx.array_family import flex
from cctbx.crystal_orientation import crystal_orientation
from cctbx.uctbx import unit_cell
from spotfinder.diffraction.imagefiles import Spotspickle_argument_module,image_files
from labelit.dptbx import AutoIndexEngine, Parameters
from labelit.command_line.imagefiles import QuickImage
from annlib_ext import AnnAdaptor

def get_precession_parameters(sources=[]):
  parameters = precession_master_phil.fetch(sources=sources)
  #validation not done here
  return parameters

def ai_factory(inputpd):
    subpd = inputpd["best_integration"]
    ai = AutoIndexEngine(inputpd["endstation"],inputpd["recommended_grid_sampling"])
    P = Parameters(xbeam=float(inputpd["xbeam"]),ybeam=float(inputpd["ybeam"]),
             distance=float(inputpd["distance"]),twotheta=float(inputpd["twotheta"]))
    ai.setBase(P)
    ai.setWavelength(float(inputpd['wavelength']))
    ai.setMaxcell(float(inputpd['ref_maxcel']))
    ai.setDeltaphi(float(inputpd['deltaphi'])*math.pi/180.)
    ai.setMosaicity(subpd["mosaicity"])
    ai.setOrientation(subpd["orient"])
    return ai

def get_col(n,sqr):
  return matrix.col((sqr[n],sqr[n+3],sqr[n+6]))

class calculate_miller:
  def __init__(self,plot_section,ori):
    labels = "HKL"
    hkl_list = plot_section.split(",")
    self.section_number = None
    for trial_plot_normal_index in [0,1,2]:
      try:
        self.section_number = int(hkl_list[trial_plot_normal_index])
        self.idx1 = (trial_plot_normal_index+1)%3
        assert hkl_list[self.idx1]==labels[self.idx1]
        self.idx2 = (trial_plot_normal_index+2)%3
        assert hkl_list[self.idx2]==labels[self.idx2]
        self.idx0 = (trial_plot_normal_index+3)%3

        break
      except:
        pass # one of three settings must succeed
    assert self.section_number!=None
    assert self.idx0!=None and self.idx1!=None and self.idx2!=None
    print self.idx0,self.idx1,self.idx2
    rmat = matrix.sqr(ori.reciprocal_matrix())
    self.unit_recip_basis1 = get_col(self.idx1,rmat).normalize() #like b*-hat
    self.unit_recip_basis2 = get_col(self.idx2,rmat).normalize() #like c*-hat
    self.a_hat_perp = self.unit_recip_basis1.cross(self.unit_recip_basis2).normalize()

    # normal to the plotted plane, like b*-hat x c*-hat
    self.c_hat_perp = self.a_hat_perp.cross(self.unit_recip_basis1)
    # second direction of plot, like a-hat-perp x b*-hat

    self.recip_basis1 = get_col(self.idx1,rmat)
    self.recip_basis2 = get_col(self.idx2,rmat)
    self.recip_basis3 = get_col(self.idx0,rmat)
    #reciprocal angle between recip_basis1 & recip_basis2
    angle = math.acos(self.unit_recip_basis1.dot(
            self.unit_recip_basis2))

    #re-expression matrix;  (f1.e1 f2.e1)
    #                       (f1.e2 f2.e2)
    self.reexpress = matrix.sqr((self.recip_basis1.length(),self.recip_basis2.length()*math.cos(angle),0.,
                          0.,                        self.recip_basis2.length()*math.sin(angle),0.,
                          0.,0.,1.)
                     ).inverse()


  def set_dmax(self,dmax):
    self.dmax=dmax

class precession_plot:
  def HKL_from_xy(self,x,y):
    # assumes x,y are page coordinates, origin at top left, x down (slow), y left to right (fast)
    # reciprocal space starts in the middle of page, axis1 to the right, axis2 up
    width = float(self.P.labelit_precession.pixel_width)/2.0
    coord1 = ((float(y) - width)/width)*(1./self.miller_calc.dmax)
    coord2 = ((-float(x) + width)/width)*(1./self.miller_calc.dmax)

    #plotting vector, represented in the orthonormal basis e = (b*-hat, c-hat-perp)
    r = matrix.col((coord1,coord2,0.0))
    #re-express in the crystal basis, f = (b*, c*)
    miller = self.miller_calc.reexpress * r

    miller3 = self.miller_calc.section_number
    hkl_list = [0.,0.,0.]
    hkl_list[self.miller_calc.idx0]=miller3
    #print hkl_list

    #hkl_zero = self.projection_onto_section_zero(hkl_list)

    hkl_list[self.miller_calc.idx1]=miller[0]
    hkl_list[self.miller_calc.idx2]=miller[1]

    #correct for the zero of this section
    hkl_list = matrix.col(hkl_list) - self.hkl_zero
    #print "hkl_list",hkl_list

    if x%50==0 and y%50==0:
      print "%6d%6d"%(x,y),"%5.1f %5.1f %5.1f"%tuple(hkl_list),
    return tuple(hkl_list)

  def projection_onto_section_zero(self,hkl):
    #compute this point, P, in orthonormal reciprocal coordinates
    P = matrix.sqr(self.rstbx_ori.reciprocal_matrix()) * matrix.col(hkl)

    # the plane normal
    N = self.miller_calc.a_hat_perp

     #projection of P onto the plane, given that a point on the plane O is 0,0,0 (otherwise Q=P-((P-O).dot(N))*N)
    Q = P - (P.dot(N))*N
    # these are still orthonormal reciprocal coordinates.  Want to convert them back to Miller indices
    M = matrix.sqr(self.rstbx_ori.direct_matrix()) * Q
    return M

  def xy_from_00(self):
    width = float(self.P.labelit_precession.pixel_width)/2.0
    miller =[self.hkl_zero[self.miller_calc.idx1],
             self.hkl_zero[self.miller_calc.idx2],
             self.hkl_zero[self.miller_calc.idx0]]

    r = self.miller_calc.reexpress.inverse()*miller
    coord1 = r[0]
    coord2 = r[1]
    x = -width*(coord2*self.miller_calc.dmax - 1.)
    y = width * (coord1*self.miller_calc.dmax + 1.)
    return x/float(self.P.labelit_precession.pixel_width),y/float(self.P.labelit_precession.pixel_width)

  def get_labelit_index_parameters(self,P):
    self.P = P # program parameters
    H = pickle.load(open("LABELIT_possible"))
    counter = [i["counter"] for i in H]
    orient = [i["orient"] for i in H]
    setting_index = counter.index(self.P.labelit_precession.bravais_choice)
    self.rstbx_ori = orient[setting_index]
    #analyze the HKL section request
    self.miller_calc = calculate_miller(self.P.labelit_precession.plot_section,self.rstbx_ori)
    miller3 = self.miller_calc.section_number
    hkl_list = [0.,0.,0.]
    hkl_list[self.miller_calc.idx0]=miller3
    self.hkl_zero = self.projection_onto_section_zero(hkl_list)

    G = pickle.load(open("LABELIT_pickle"))

    w = float( G['wavelength'] )
    self.ai = ai_factory(G)

    self.ai.setOrientation(self.rstbx_ori)
    self.pixelsz = float(G["pixel_size"])
    template = G["file"][G["file"].keys()[0]]
    I = image_files(Spotspickle_argument_module(template))
    a_file = I.filenames.FN[0]
    self.template = os.path.join(a_file.cwd,a_file.template)
    self.wildcard_count = self.template.count("#")

    return self

  def construct_image_lookup(self):
    self.reverse_lookup = {}
    self.all_image_centers = flex.double()
    self.deltaphi = None
    for x in xrange(self.P.labelit_precession.image_range[0],
      self.P.labelit_precession.image_range[1]+1):
      a_file = self.template.replace("#"*self.wildcard_count, ("%%0%sd"%self.wildcard_count)%x)
      Q = QuickImage(a_file)
      if self.deltaphi == None:  self.deltaphi = Q.deltaphi
      assert self.deltaphi == Q.deltaphi
      phivalue = (Q.osc_start + (Q.deltaphi/2.0))%360.0
      #print "image",x,"phi",phivalue
      self.reverse_lookup[phivalue]=x
      self.all_image_centers.append(phivalue)
    self.distance_tree = AnnAdaptor(self.all_image_centers,1)

  def imagefunc(self,deg):
    modular_deg = deg%360.
    self.distance_tree.query(flex.double([modular_deg]))
    if math.sqrt(self.distance_tree.distances[0]) > self.deltaphi/2.0: return None
    image_number = self.reverse_lookup.get(self.all_image_centers[self.distance_tree.nn[0]])
    return image_number

  def get_plot_values(self):
    param = self.P
    all_images={} # a very large memory sink.  Might need a more clever implementation later.
                  # relies on having all the images in memory at the same time.
    init_value = 235
    self.all_values = flex.int(param.labelit_precession.pixel_width**2,init_value)
    self.all_values.resize(flex.grid(param.labelit_precession.pixel_width,
                                     param.labelit_precession.pixel_width))
    scale_factor = 255./param.labelit_precession.intensity_full_scale
    uc = self.rstbx_ori.unit_cell()
    self.miller_calc.set_dmax(param.labelit_precession.resolution_outer)

    for x in xrange(param.labelit_precession.pixel_width):
      for y in xrange(param.labelit_precession.pixel_width):
        #print x,y
        H,K,L = self.HKL_from_xy(x,y)

        resolution = uc.d((int(H),int(K),int(L)))
        if x%50==0 and y%50==0:
          print "%6.1f"%resolution
        if resolution < param.labelit_precession.resolution_outer :
          continue

        for item in self.ai.predict_hkl((H,K,L),self.miller_calc.dmax):
          #this is a kludge!  if the HKL is recorded through the Ewald sphere twice,
          # only pick up the first time.
          deg  = item[2]*180./math.pi
          image_number = self.imagefunc(deg)
          if image_number == None: continue
          pixelx = int(item[0]/self.pixelsz)
          pixely = int(item[1]/self.pixelsz)

          if not all_images.has_key(image_number):
            filenm = self.template.replace("#"*self.wildcard_count, ("%%0%sd"%self.wildcard_count)%image_number)
            Q = QuickImage(filenm)
            Q.read()
            all_images[image_number]=Q
          this_data = all_images[image_number].linearintdata
          pxl_value = this_data[pixelx,pixely]
          #print "%4d %4d %7.3f%7.3f%7.3f,%7.3f,%7.2f->%3d x%4d y%4d val %6d"%(
          #  x,y,H,K,L,resolution,deg,image_number,pixelx,pixely,pxl_value)


          plotted_value = int(pxl_value * scale_factor)
          if plotted_value >  255:
            plotted_value = 255
          if plotted_value <  0:
            plotted_value = 0

          self.all_values[x,y] = 255 - plotted_value
          break

from reportlab.pdfgen.canvas import Canvas
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch,cm,mm

page_origin = (20.,190.)
boxedge = 500.

class Graph:

  def __init__(self,fileout):
    self.c = Canvas(fileout,pagesize=letter)

  def title(self,text):
    lines = text.split('\n')
    self.c.setFont('Helvetica',12)
    for x in xrange(len(lines)):
      self.c.drawString(2*cm,(26-0.5*x)*cm,lines[x])

  def image(self,pi):
    self.c.drawInlineImage(pi,page_origin[0],page_origin[1],
      width=boxedge,height=boxedge)
    self.c.setFillColorRGB(1,0,0)
  def dot(self,dc):
    self.c.circle(page_origin[0]+(dc[1]*boxedge),page_origin[1]+(boxedge-dc[0]*boxedge),r=1.,stroke=0,fill=1)

  def __del__(self):
    self.c.save()

class GeneratePDF:
  def __init__(self,project):
    self.image_size = project.P.labelit_precession.pixel_width
    self.R = Graph(project.P.labelit_precession.pdf_output.file)
    self.data = project.all_values
    self.section = project.P.labelit_precession.plot_section
    self.project = project

  def make_image_plots_detail(self):
    import Image
    mode = "L" # black and white
    HWsizes = (self.image_size, self.image_size)
    pil_image = Image.new(mode,HWsizes)

    data = list(self.data)

    pil_image.putdata(data)

    uc = self.project.rstbx_ori.unit_cell()
    axis = ["a*","b*","c*"][self.project.miller_calc.idx1]
    axis2 = ["a*","b*","c*"][self.project.miller_calc.idx2]
    angle = math.acos(self.project.miller_calc.unit_recip_basis1.dot(
            self.project.miller_calc.unit_recip_basis2))*180./math.pi


    self.R.title("One reciprocal lattice plane, %s. Red spot is the %s"%(self.section,
    self.section.replace("H","0").replace("K","0").replace("L","0"))+"""
Unit cell: %s"""%(uc)+"""
%s points in the --> direction starting at the red spot"""%axis+"""
%s is in the plane at a %5.1f degree angle"""%(axis2,angle) )

    self.R.image(pil_image)
    dotcoords = self.project.xy_from_00()
    self.R.dot(dotcoords)

    self.R.c.showPage()
    return self

def run(args):
  if args==[]:
    print usage()
    return
  argument_interpreter = libtbx.phil.command_line.argument_interpreter(
    master_phil=precession_master_phil,home_scope="hidden_symmetry")
  parsed_arguments = []
  for arg in args:
    if arg.find("help")>=0:
      param_master = get_precession_parameters(sources=[])
      param_master.show(attributes_level=1)
      return
    else:
      try:
        a = argument_interpreter.process(arg=arg)
      except:
        print "Unknown file or keyword: %s"%arg
        print
        return
      else:
        parsed_arguments.append(a)
  param_master = get_precession_parameters(sources=parsed_arguments)
  param_master.show()
  param = param_master.extract()
  PP = precession_plot().get_labelit_index_parameters(param)
  PP.construct_image_lookup()
  PP.get_plot_values()
  PDF = GeneratePDF(PP)
  PDF.make_image_plots_detail()

if __name__=="__main__":
  import sys
  run(sys.argv[1:])