from phaser.pipeline import process_model_new as process_model
from phaser.pipeline import mr_object_new as mr_object
from phaser.pipeline import phaser_ai_new as phaser_ai

import phaser
import libtbx.load_env
from iotbx import bioinformatics

import os.path
import sys
import logging

def read_data(file_name, fp, sigfp, hires = None):
    
    inp = phaser.InputMR_DAT()
    inp.setHKLI( file_name )
    inp.setLABI( fp, sigfp )
    inp.setMUTE( True )
    
    if hires:
        inp.setHIRES( hires )
    
    result = phaser.runMR_DAT( inp )
    
    if result.Failed():
        raise Exception, result
    
    symm = result.getSpaceGroupHall()
    cell = result.getUnitCell()
    inp = phaser.InputANO()
    inp.setREFL( result.getMiller(), result.getF(), result.getSIGF() )
    inp.setCELL6( cell )
    inp.setSPAC_HALL( symm )
    inp.setMUTE( True )
    inp.setHKLO( False )
    
    result = phaser.runANO( inp )
    
    if result.Failed():
        raise Exception, result
    
    xray_data = mr_object.XrayData(
        miller = result.getMiller(),
        fp = result.getCorrectedF(),
        sigfp = result.getCorrectedSIGF()
        )
    spac = mr_object.SpaceGroup( space_group_hall = symm )
    
    unit_cell = mr_object.UnitCell( unit_cell = cell )
    
    return ( xray_data, spac, unit_cell )


def lysozyme_from_models(space_groups):
    
    root = libtbx.env.under_dist( "phaser", "tutorial" )
    mtz_file = os.path.join( root, "iod_scala-unique.mtz" )
    pdb_file1 = "/home/gaborb/data/lyso2001/193l_prot.pdb"
    pdb_file2 = os.path.join( root, "1fkq_prot.pdb" )
    
    ( xray_data, spac, cell ) = read_data( mtz_file, "F_New", "SIGF_New", 3.5 )
    hewl_seq = bioinformatics.sequence(
        name = "HEWL",
        sequence = ( "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQIN"
            + "SRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTD"
            + "VQAWIRGCRL" )
        )
    component = mr_object.Component( mw = 17000, sequence = hewl_seq )
    composition = mr_object.Composition( components = [ component ] )
    mr_ensemble1 = process_model.SimpleEnsemble(
        pdb_file = pdb_file1,
        identity = 0.99,
        component = component
        )
    mr_ensemble2 = process_model.SimpleEnsemble(
        pdb_file = pdb_file2,
        identity = 0.45,
        component = component
        )
    mr_ensemble3 = process_model.SimpleEnsemble(
        pdb_file = "/home/gaborb/data/lyso/1jhl_prot.pdb",
        identity = 0.92,
        component = component
        )
    mr_ensemble4 = process_model.SimpleEnsemble(
        pdb_file = "/home/gaborb/data/lyso/2gv0_prot.pdb",
        identity = 0.68,
        component = component
        )
    mr_ensemble5 = process_model.SimpleEnsemble(
        pdb_file = "/home/gaborb/data/lyso/1gf3_prot.pdb",
        identity = 0.61,
        component = component
        )
    mr_ensemble6 = process_model.SimpleEnsemble(
        pdb_file = "/home/gaborb/data/lyso/2z2e_prot.pdb",
        identity = 0.50,
        component = component
        )
    mr_ensemble7 = process_model.SimpleEnsemble(
        pdb_file = "/home/gaborb/data/lyso/2h5z_prot.pdb",
        identity = 0.40,
        component = component
        )
    ensembles = [ mr_ensemble1, mr_ensemble2, mr_ensemble3, mr_ensemble4,
        mr_ensemble5, mr_ensemble6, mr_ensemble7 ]
    problem = mr_object.Problem(
        xray_data = xray_data,
        unit_cell = cell
        )
    ai = phaser_ai.PhaserAI.MRMode(
        problem = problem,
        composition = composition,
        symmetries = space_groups,
        outstream = sys.stdout
        )
    
    for ense in ensembles:
        ai.add( bbobj = ense )
    
    ai.logger().setLevel( logging.INFO )
    
    ai.solve()
    
    
def lysozyme_from_hits():
    
    root = libtbx.env.under_dist( "phaser", "tutorial" )
    mtz_file = os.path.join( root, "iod_scala-unique.mtz" )
    pdb_file1 = "/home/gaborb/data/lyso2001/193l_prot.pdb"
    pdb_file2 = os.path.join( root, "1fkq_prot.pdb" )
    
    ( xray_data, spac, cell ) = read_data( mtz_file, "F_New", "SIGF_New", 3.5 )
    hewl_seq = bioinformatics.sequence(
        name = "HEWL",
        sequence = ( "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQIN"
            + "SRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTD"
            + "VQAWIRGCRL" )
        )
    lact_seq = bioinformatics.sequence(
        name = "LACT",
        sequence = ( "MEQLTKCEVFQKLKDLKDYGGVSLPEWVCVAFHTSGYDTQAIVQNNDSTEYGLFQIN"
            + "NKIWCKDDQNPHSRNICNISCDKFLDDDLTDDIVCAKKILDKVGINYWLAHKALCSEK"
            + "LDQWLCEKL" )
        )
    component = mr_object.Component( mw = 17000, sequence = hewl_seq )
    composition = mr_object.Composition( components = [ component ] )
    ( ali, bad ) = bioinformatics.ali_alignment_parse(
        file( "ali_193l_193l.fa" ).read()
        )
    assert ali
    b1 = process_model.BlastHit(
        identifier = "193L_A",
        component = component,
        score = 0.9
        )
    b1.set_sequence( hewl_seq )
    b1.set_structure(
        mr_object.Structure( file_name = os.path.join( "pipeline", "193L_A.pdb" ) )
        )
    a1 = process_model.Alignment(
        blast_hit = b1,
        alignment = ali
        )
    ( ali, bad ) = bioinformatics.ali_alignment_parse(
        file( "ali_193l_1fkq.fa" ).read()
        )
    assert ali
    b2 = process_model.BlastHit(
        identifier = "1FKQ_A",
        component = component,
        score = 0.5
        )
    b2.set_sequence( lact_seq )
    b2.set_structure(
        mr_object.Structure( file_name = os.path.join( "pipeline", "1FKQ_A.pdb" ) )
        )
    a2 = process_model.Alignment(
        blast_hit = b2,
        alignment = ali
        )
    problem = mr_object.Problem(
        xray_data = xray_data,
        unit_cell = cell
        )
    ai = phaser_ai.PhaserAI(
        problem = problem,
        composition = composition,
        symmetries = [ spac ],
        extension = phaser_ai.SingleExtension.Autofetch()
        )
    
    ai.add( bbobj = a1 )
    ai.add( bbobj = b1 )
    ai.add( bbobj = a2 )
    ai.add( bbobj = b2 )
    
    ai.solve()
    
    
if __name__ == "__main__":
    sgs = [
        mr_object.SpaceGroup( space_group_hall = " P 4nw 2abw" ),
        #mr_object.SpaceGroup( space_group_hall = " P 4abw 2nw" ),
        ]
    lysozyme_from_models( space_groups = sgs )