#!/usr/bin/env python
# MosflmCheckIndexerSolution.py
#   Copyright (C) 2009 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
# 12th May 2009
#
# Code to check the autoindex solution from mosflm for being
# pseudo-centred (i.e. comes out as centered when it should not be)
#

import os
import math
import sys

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

if not os.environ['XIA2_ROOT'] in sys.path:
    sys.path.append(os.environ['XIA2_ROOT'])

# xia2 stuff...

from Wrappers.XIA.Printpeaks import Printpeaks
from Wrappers.XIA.Diffdump import Diffdump
from Handlers.Streams import Debug
from Handlers.Flags import Flags
from lib.bits import nint
from Experts.MatrixExpert import format_matrix
from Experts.LatticeExpert import l2s, s2l

# optional labelit stuff

from Wrappers.Labelit.DistlSignalStrength import DistlSignalStrength

# things we can work on...

from Wrappers.CCP4.MosflmWrapper import MosflmWrapper

use_distl = False

# cctbx stuff

from cctbx import sgtbx
from cctbx import crystal
from cctbx import uctbx
from scitbx import matrix

# check for deprecation, add workaround (thanks to RWGK 21/APR/10)

if (hasattr(matrix.rec, "rotate_around_origin")):
    matrix.rec.rotate = matrix.rec.rotate_around_origin

# end workaround

# failover on spot picking...

def locate_maxima(image):

    global use_distl

    if use_distl:
        try:
            dss = DistlSignalStrength()
            dss.set_image(image)
            peaks = dss.find_peaks()

            if not peaks:
                raise RuntimeError, 'no peaks found'

        except Exception, e:
            use_distl = False

    if not use_distl:

        pp = Printpeaks()
        pp.set_image(image)
        peaks = pp.get_maxima()

    return peaks

def mosflm_check_indexer_solution(indexer):

    distance = indexer.get_indexer_distance()
    axis = matrix.col([0, 0, 1])
    beam = indexer.get_indexer_beam()
    cell = indexer.get_indexer_cell()
    wavelength = indexer.get_wavelength()

    space_group_number = l2s(indexer.get_indexer_lattice())
    spacegroup = sgtbx.space_group_symbols(space_group_number).hall()
    phi = indexer.get_header()['phi_width']

    sg = sgtbx.space_group(spacegroup)

    if not (sg.n_ltr() - 1):
        # primitive solution - just return ... something
        return None, None, None, None

    # FIXME need to raise an exception if this is not available!
    m_matrix = indexer.get_indexer_payload('mosflm_orientation_matrix')

    # N.B. in the calculation below I am using the Cambridge frame
    # and Mosflm definitions of X & Y...

    m_elems = []

    for record in m_matrix[:3]:
        record = record.replace('-', ' -')
        for e in map(float, record.split()):
            m_elems.append(e / wavelength)

    mi = matrix.sqr(m_elems)
    m = mi.inverse()

    A = matrix.col(m.elems[0:3])
    B = matrix.col(m.elems[3:6])
    C = matrix.col(m.elems[6:9])

    # now select the images - start with the images that the indexer
    # used for indexing, though can interrogate the FrameProcessor
    # interface of the indexer to put together a completely different
    # list if I like...

    images = []

    for i in indexer.get_indexer_images():
        for j in i:
            if not j in images:
                images.append(j)

    images.sort()

    # now construct the reciprocal-space peak list n.b. should
    # really run this in parallel...

    spots_r = []

    spots_r_j =  { }

    for i in images:
        image = indexer.get_image_name(i)
        dd = Diffdump()
        dd.set_image(image)
        header = dd.readheader()
        phi = header['phi_start'] + 0.5 * header['phi_width']
        pixel = header['pixel']
        wavelength = header['wavelength']
        peaks = locate_maxima(image)

        spots_r_j[i] = []

        for p in peaks:
            x, y, isigma = p

            if isigma < 5.0:
                continue

            xp = pixel[0] * y - beam[0]
            yp = pixel[1] * x - beam[1]

            scale = wavelength * math.sqrt(
                xp * xp + yp * yp + distance * distance)

            X = distance / scale
            X -= 1.0 / wavelength
            Y = - xp / scale
            Z = yp / scale

            S = matrix.col([X, Y, Z])

            rtod = 180.0 / math.pi

            spots_r.append(S.rotate(axis, - phi / rtod))
            spots_r_j[i].append(S.rotate(axis, - phi / rtod))

    # now reindex the reciprocal space spot list and count - n.b. need
    # to transform the Bravais lattice to an assumed spacegroup and hence
    # to a cctbx spacegroup!

    # lists = [spots_r_j[j] for j in spots_r_j]
    lists = []
    lists.append(spots_r)

    for l in lists:

        absent = 0
        present = 0
        total = 0

        for spot in l:
            hkl = (m * spot).elems

            total += 1

            ihkl = map(nint, hkl)

            if math.fabs(hkl[0] - ihkl[0]) > 0.1:
                continue

            if math.fabs(hkl[1] - ihkl[1]) > 0.1:
                continue

            if math.fabs(hkl[2] - ihkl[2]) > 0.1:
                continue

            # now determine if it is absent

            if sg.is_sys_absent(ihkl):
                absent += 1
            else:
                present += 1

        # now perform the analysis on these numbers...

        sd = math.sqrt(absent)

        if total:

            Debug.write('Counts: %d %d %d %.3f' % \
                        (total, present, absent, (absent - 3 * sd) / total))

        else:

            Debug.write('Not enough spots found for analysis')
            return False, None, None, None

        if (absent - 3 * sd) / total < 0.008:
            return False, None, None, None

    # in here need to calculate the new orientation matrix for the
    # primitive basis and reconfigure the indexer - somehow...

    # ok, so the bases are fine, but what I will want to do is reorder them
    # to give the best primitive choice of unit cell...

    sgp = sg.build_derived_group(True, False)
    lattice_p = s2l(sgp.type().number())
    symm = crystal.symmetry(unit_cell = cell,
                            space_group = sgp)

    rdx = symm.change_of_basis_op_to_best_cell()
    symm_new = symm.change_basis(rdx)

    # now apply this to the reciprocal-space orientation matrix mi

    # cb_op = sgtbx.change_of_basis_op(rdx)
    cb_op = rdx
    R = cb_op.c_inv().r().as_rational().as_float().transpose().inverse()
    mi_r = mi * R

    # now re-derive the cell constants, just to be sure

    m_r = mi_r.inverse()
    Ar = matrix.col(m_r.elems[0:3])
    Br = matrix.col(m_r.elems[3:6])
    Cr = matrix.col(m_r.elems[6:9])

    a = math.sqrt(Ar.dot())
    b = math.sqrt(Br.dot())
    c = math.sqrt(Cr.dot())

    rtod = 180.0 / math.pi

    alpha = rtod * Br.angle(Cr)
    beta = rtod * Cr.angle(Ar)
    gamma = rtod * Ar.angle(Br)

    # print '%6.3f %6.3f %6.3f %6.3f %6.3f %6.3f' % \
    # (a, b, c, alpha, beta, gamma)

    cell = uctbx.unit_cell((a, b, c, alpha, beta, gamma))

    amat = [wavelength * e for e in mi_r.elems]
    bmat = matrix.sqr(cell.fractionalization_matrix())
    umat = mi_r * bmat.inverse()

    # yuk! surely I don't need to do this...

    # I do need to do this, and don't call me shirley!

    new_matrix = ['%s\n' % r for r in \
                  format_matrix((a, b, c, alpha, beta, gamma),
                                amat, umat.elems).split('\n')]

    # ok - this gives back the right matrix in the right setting - excellent!
    # now need to apply this back at base to the results of the indexer.

    # N.B. same should be applied to the same calculations for the XDS
    # version of this.

    return True, lattice_p, new_matrix, (a, b, c, alpha, beta, gamma)

if __name__ == '__main__':

    # run a test!

    from Modules.Indexer.IndexerFactory import Indexer

    i = Indexer()

    i.setup_from_image(sys.argv[1])

    print 'Refined beam is: %6.2f %6.2f' % i.get_indexer_beam()
    print 'Distance:        %6.2f' % i.get_indexer_distance()
    print 'Cell: %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f' % i.get_indexer_cell()
    print 'Lattice: %s' % i.get_indexer_lattice()
    print 'Mosaic: %6.2f' % i.get_indexer_mosaic()

    status = mosflm_check_indexer_solution(i)

    if status is True:
        print 'putative centred solution came out as wrong'

    elif status is False:
        print 'putative centred solution came out as right'

    elif status is None:
        print 'putative solution not centred'