#!/usr/bin/env python

from phaser import mr_solution, mr_ndim, mr_rlist, mr_set, MapCoefs, map_str_float1D
from scitbx.math import euler_angles
from cctbx.xray import scatterer
from cctbx.array_family import flex

from cStringIO import StringIO
import re
import math

class LabinParser(dict):
    """
    CCP4-style labin parsing
    """

    @staticmethod
    def parse(labin):

        obj = LabinParser()

        for labin_definition in labin.split():
            ( program_label, data_label ) = labin_definition.split( "=" )
            obj[ program_label ] = data_label

        return obj


class SolutionFileParser(object):
    """
    This class parses a solution file and creates an mr_solution object
    """

    SOLU_TRIAL = re.compile(
        r"""
        ^ SOLU \w* \s+ TRIAL \s+ ENSE \w* \s+
        (\w+)
        \s+ EULER
        \s+ ( -? \d+ \. \d+ )
        \s+ ( -? \d+ \. \d+ )
        \s+ ( -? \d+ \. \d+ )
        \s+ RFZ \s+
        ( \d+ \. \d+ )
        """,
        re.VERBOSE
        )

    SOLU_6DIM = re.compile(
        r"""
        ^ SOLU \w* \s+ 6DIM \s+ ENSE \w* \s+
        ([\w\{\}\-]+)
        \s+ EULER
        ( \s+ -?\d+ \. \d+ )
        ( \s+ -?\d+ \. \d+ )
        ( \s+ -?\d+ \. \d+ )
        \s+
        ( FRAC | ORTH )
        ( \s+ -?\d+ \. \d+ )
        ( \s+ -?\d+ \. \d+ )
        ( \s+ -?\d+ \. \d+ )
        """,
        re.VERBOSE
        )

    SOLU_BFAC = re.compile(
        r"""
        \s+ BFAC \s+ ( -?\d+ \. \d+ )
        """,
        re.VERBOSE
        )

    SOLU_SET = re.compile(
        r"""
        ^ SOLU \w* \s+ SET
        ( [^\n]* )
        """,
        re.VERBOSE
        )

    SOLU_SPAC = re.compile(
        r"""
        ^ SOLU \w* \s+ SPAC
        ( [^\n]* )
        """,
        re.VERBOSE
        )

    SOLU_FSOL = re.compile(
        r"""
        ^ SOLU \w* \s+ SOLPAR \s+ FSOL \s+ ( -?\d+ \. \d+ )
        ( [^\n]* )
        """,
        re.VERBOSE
        )

    SOLU_BSOL = re.compile(
        r"""
        ^ SOLU \w* \s+ SOLPAR \s+ .* \s+ BSOL \s+ ( -?\d+ \. \d+ )
        """,
        re.VERBOSE
        )

    SOLU_FIXB = re.compile(
        r"""
        ^ SOLU \w* \s+ .* \s+ FIXB \s+ ( ON | OFF )
        """,
        re.VERBOSE
        )

    SOLU_VRMS_1 = re.compile(
        r"""
        ^ SOLU \w* \s+ ENSE \w* \s+
        ( \S+ ) \s+
        VRMS
        ( (?: \s+ -? \d+ \. \d+ )+ )
        """,
        re.VERBOSE
        )

    SOLU_VRMS_2 = re.compile(
        r"""
        ^ SOLU \w* \s+ ENSE \w* \s+
        ( \S+ ) \s+
        VRMS \s DELTA
        ( \s+ [+-] \d+ \. \d+ ) \s+
        RMSD
        ( (?: \s+ -? \d+ \. \d+ )+ )
        """,
        re.VERBOSE
        )

    SOLU_VRMS_3 = re.compile(
        r"""
        ^ SOLU \w* \s+ ENSE \w* \s+
        ( \S+ ) \s+
        VRMS \s DELTA
        ( \s+ [+-] \d+ \. \d+ ) \s+
        RMSD
        ( (?: \s+ -? \d+ \. \d+ )+ ) \s+
        \# VRMS
        """,
        re.VERBOSE
        )

    def __init__(self, file):

        self.mr_sets = []
        self.empty_current()

        for line in file:
            self.parse_sol_line( line )

        self.move_current_to_processed()


    def parse_sol_line(self, line):

        result = SolutionFileParser.SOLU_TRIAL.search( line )

        if result:
            self.current_rlist.append( mr_rlist(
                    result.group(1),
                    (
                            float( result.group( 2 ) ),
                            float( result.group( 3 ) ),
                            float( result.group( 4 ) ),
                    ),
                    0,
                    float( result.group( 5 ) ) )
            )
            return

        result = SolutionFileParser.SOLU_6DIM.search( line )

        if result:
            bfac_result =  SolutionFileParser.SOLU_BFAC.search( line )
            self.current_ndim.append( mr_ndim(
                result.group(1),
                euler_angles.zyz_matrix(
                    float( result.group( 2 ) ),
                    float( result.group( 3 ) ),
                    float( result.group( 4 ) ),
                ),
                ( result.group( 5 ) == "FRAC" ),
                ( result.group( 5 ) == "FRAC" ),
                (
                    float( result.group( 6 ) ),
                    float( result.group( 7 ) ),
                    float( result.group( 8 ) ),
                ),
                float( bfac_result.group( 1 ) ) if bfac_result else 0,
                ( line.find( "FIXR" ) != -1 ),
                ( line.find( "FIXT" ) != -1 ),
                ( line.find( "FIXB" ) != -1 ),
                ) )

            return

        result = SolutionFileParser.SOLU_SET.search( line )

        if result:
            if self.current_annotation:
                self.move_current_to_processed()
                self.empty_current()
            self.current_annotation = result.group( 1 )
            return

        try:
            ( name, vrms ) = self.get_vrms_value( line = line )
            self.current_vrms[ name ] = vrms
            return

        except SyntaxError:
            pass

        result = SolutionFileParser.SOLU_SPAC.search( line )

        if result:
            self.current_space_group = result.group( 1 ).strip()

        result = SolutionFileParser.SOLU_FSOL.search( line )

        if result:
            self.current_fsol = float(result.group(1))

        result = SolutionFileParser.SOLU_BSOL.search( line )

        if result:
            self.current_bsol = float(result.group(1))

        result = SolutionFileParser.SOLU_FIXB.search( line )

        if result:
            self.current_fixb = ( True if result.group(1) == "ON" else False )



    def empty_current(self):

        self.current_annotation = None
        self.current_rlist = []
        self.current_ndim = []
        self.current_vrms = map_str_float1D()
        self.current_space_group = ""
        self.current_fsol = 0
        self.current_bsol = 0
        self.current_fixb = False


    def move_current_to_processed(self):

        from cctbx import sgtbx

        if self.current_space_group:
            sg = sgtbx.space_group_info( symbol = self.current_space_group ).type().hall_symbol()

        else:
            sg = ""

        self.mr_sets.append( mr_set(
                self.current_annotation,
                self.current_ndim,
                self.current_rlist,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                [],
                MapCoefs(),
                self.current_vrms,
                sg,
                0,
                ) )


    def get_dot_sol(self):

        return mr_solution( self.mr_sets )


    @staticmethod
    def parse_file( file_name ):

        return SolutionFileParser( file( file_name ) )


    @staticmethod
    def parse_string( sol_string ):

        return SolutionFileParser( StringIO( sol_string ) )


    @classmethod
    def get_vrms_value(cls, line):

        result = cls.SOLU_VRMS_1.search( line )

        if result:
            ( name, rms ) = result.groups()
            return ( name, [ float( r ) for r in rms.split() ] )

        result = cls.SOLU_VRMS_3.search( line )

        if result:
            ( name, d, rms ) = result.groups()
            delta = float( d.strip() )
            return ( name, [ math.sqrt( float( r ) ** 2 + delta ) for r in rms.split() ] )

        result = cls.SOLU_VRMS_2.search( line )

        if result:
            ( name, d, rms ) = result.groups()
            return ( name, [ float( r ) for r in rms.split() ] )

        raise SyntaxError, "Not found"


class AtomDefinition(object):
    """
    Information about an atom
    """

    def __init__(self, xtal, elem, x, y, z, occ, bfac):

        self.xtal = xtal
        self.elem = elem
        self.x = x
        self.y = y
        self.z = z
        self.occ = occ
        self.bfac = bfac


class AtomFileParser(object):
    """
    This class parses an EP .sol-file and return atom records
    """

    ATOM_CRYS = re.compile(
            r"""
            ^ \s* ATOM \s+ CRYS \w* \s+
            (\w+)
            \s+ TYPE \s+
            (\w+)
            \s+ ( FRAC | ORTH )
            ( \s+ -?\d+ \. \d* )
            ( \s+ -?\d+ \. \d* )
            ( \s+ -?\d+ \. \d* )
            \s+ OCC
            ( \s+ -?\d+ \. \d* )
            """,
            re.VERBOSE )

    ISOB_BFAC = re.compile(
            r"""
            \s+ ISOB
            ( \s+ -?\d+ \. \d* )
            """,
            re.VERBOSE )

    F_PRIME_AND_F_DOUBLEPRIME = re.compile(
            r"""
            ^ \s* SCAT \w* \s+ TYPE \s+
            ( \w+ )
            \s+ FP \s* = \s*
            ( -?\d+ \. \d+ )
            \s+ FDP \s* = \s*
            ( -?\d+ \. \d+ )
            """,
            re.VERBOSE )

    def __init__(self, stream):

        self.atoms = []
        self.form_factors = {}

        for line in stream:
            self.parse_sol_line( line )


    def parse_sol_line(self, line):

        result = AtomFileParser.ATOM_CRYS.search( line )

        if result:
            bfac_search = AtomFileParser.ISOB_BFAC.search( line )

            if bfac_search:
                bfac = float( bfac_search.group( 1 ) )

            else:
                bfac = None

            self.atoms.append(
                    AtomDefinition(
                            result.group(1),
                            result.group( 2 ),
                            float( result.group( 4 ) ),
                            float( result.group( 5 ) ),
                            float( result.group( 6 ) ),
                            float( result.group( 7 ) ),
                            bfac )
                )
            return

        result = AtomFileParser.F_PRIME_AND_F_DOUBLEPRIME.search( line )

        if result:
            self.form_factors[ result.group( 1 ) ] = (
                    float( result.group( 2 ) ),
                    float( result.group( 3 ) )
                )
            return


    def get_f_p(self, element):

        return self.form_factors.get( element, ( 0.0, 0.0 ) )[ 0 ]


    def get_f_dp(self, element):

        return self.form_factors.get( element, ( 0.0, 0.0 ) )[ 1 ]


    def get_cctbx_atoms(self):

        return flex.xray_scatterer(
            [ scatterer(
                    scattering_type = atom.elem,
                    site = ( atom.x, atom.y, atom.z ),
                    occupancy = atom.occ,
                    b = atom.bfac if atom.bfac else -8.0 * math.pi * math.pi
                    )
                for atom in self.atoms ]
            )


    @staticmethod
    def get_empty_atoms_list(self):
        return flex.xray_scatterer( [] )


    @staticmethod
    def parse_string( atom_string ):

        return AtomFileParser( StringIO( atom_string ) )


    @staticmethod
    def parse_file( file_name ):

        return AtomFileParser( file( file_name ) )