from phaser.pipeline import phaser_ai
from phaser.pipeline import knowledge_sources
from phaser.pipeline import process_model
from phaser.pipeline import openbba
from phaser.pipeline import event
from phaser.pipeline import strategy
from phaser.pipeline import engine
from phaser.pipeline.component import analyser

from phaser import pickle_support

import phaser

import sys
import os.path

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
    
    inp = phaser.InputANO()
    inp.setREFL( result.getMiller(), result.getF(), result.getSIGF() )
    inp.setCELL6( result.getUnitCell() )
    inp.setSPAC_HALL( result.getSpaceGroupHall() )
    inp.setMUTE( True )
    inp.setHKLO( False )
    
    result = phaser.runANO( inp )
    
    if result.Failed():
        raise Exception, result
    
    xray_data = process_model.XrayData(
        miller = result.getMiller(),
        fp = result.getCorrectedF(),
        sigfp = result.getCorrectedSIGF()
        )
    symmetry = process_model.Symmetry(
        space_group_hall = result.getSpaceGroupHall(),
        unit_cell = result.getUnitCell()
        )
    
    return ( xray_data, symmetry )


def get_mr_solver(selection_strategy, engine, analysers = []):
    
    return phaser_ai.Controller(
        initiators = [
            knowledge_sources.ExtensionStanceFind(
                mr_set_selection_strategy = selection_strategy
                ),
            ],
        processors = [
            knowledge_sources.ExtensionResultAnalyse(
                optional_analysers = analysers
                ),
            knowledge_sources.ExtensionCalculationProcess(),
            knowledge_sources.ExtensionCalculationExecute( engine = engine ),
            knowledge_sources.RegularExtensionCalculationSetup(),
            ]
        )


def setup_solution_blackboard(startup_objects):
    
    blackboard = openbba.Blackboard( memory = openbba.memory.Transient() )
    
    for bbobj in startup_objects:
        bbobj.attach_to( blackboard = blackboard )
        
    return blackboard


def get_loggers():
    
    return [
        phaser_ai.StreamLogger(
            event_type = event.Message,
            display = lambda event: str( event ),
            stream = sys.stdout
            ),
        phaser_ai.StreamLogger(
            event_type = event.CalculationComplete,
            display = lambda event: str( event ),
            stream = sys.stdout
            ),
        phaser_ai.StreamLogger(
            event_type = event.Controller,
            display = lambda event: str( event ),
            stream = sys.stdout
            ),
        ]
    


def lysozyme():
    
    root = "/home/gaborb/workspace/phaser/src/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, symmetry ) = read_data( mtz_file, "F_New", "SIGF_New", 3.5 )
    component = process_model.Component(
        sequence = ( "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINS"
            + "RWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDV"
            + "QAWIRGCRL" )
        )
    composition = process_model.Composition.FromComponents( components = [ component ] )
    mr_model1 = process_model.MRModel.Identity( pdb_file = pdb_file1, value = 0.90 )
    mr_ensemble1 = process_model.MREnsemble(
        component = component,
        models = [ mr_model1 ],
        score = 0.90
        )
    mr_model2 = process_model.MRModel.Identity( pdb_file = pdb_file2, value = 0.450 )
    mr_ensemble2 = process_model.MREnsemble(
        component = component,
        models = [ mr_model2 ],
        score = 0.45
        )
    problem = process_model.MRProblem(
        xray_data = xray_data,
        symmetry = symmetry,
        composition = composition
        )
    
    solver = get_mr_solver(
        selection_strategy = strategy.mr_set.LocalTop(),
        engine = engine.Multiprocessing(),
        analysers = [
            analyser.SolutionIdentify(),
            analyser.SearchThreadTerminate()
            ]
        )
    cpu1 = process_model.CPU()
    cpu2 = process_model.CPU()
    loggers = get_loggers()
    blackboard = setup_solution_blackboard(
        startup_objects = [ problem, component, mr_ensemble1, mr_ensemble2, cpu1, cpu2 ]
        )
    
    for logger in loggers:
        blackboard.register( logger )
        
    solver.run( blackboard = blackboard )
    

def a31p():
    
    root = "/home/gaborb/workspace/phaser/src/tutorial"
    mtz_file = os.path.join( root, "A31P_3A.mtz" )
    pdb_file = os.path.join( root, "helix.pdb" )
    
    ( xray_data, symmetry ) = read_data( mtz_file, "FP", "SIGFP" )
    component = process_model.Component(
        sequence = "A" * 26
        )
    composition = process_model.Composition.FromComponents( components = [ component ] * 4 )
    mr_model = process_model.MRModel.Rms( pdb_file = pdb_file, value = 1.0 )
    mr_ensemble = process_model.MREnsemble(
        component = component,
        models = [ mr_model ],
        score = 0.45
        )
    problem = process_model.MRProblem(
        xray_data = xray_data,
        symmetry = symmetry,
        composition = composition
        )
    solver = get_mr_solver(
        selection_strategy = strategy.mr_set.FrozenGlobalTop(),
        engine = engine.Multiprocessing(),
        analysers = [
            analyser.SolutionIdentify(),
            analyser.SearchThreadTerminate()
            ]
        )
    cpu1 = process_model.CPU()
    cpu2 = process_model.CPU()
    loggers = get_loggers()
    blackboard = setup_solution_blackboard(
        startup_objects = [ problem, component, mr_ensemble, cpu1, cpu2 ]
        )
    
    for logger in loggers:
        blackboard.register( logger )
        
    solver.run( blackboard = blackboard )
    
    
def slt_i():
    
    root = "/home/gaborb/data/slt_i"
    mtz_file = os.path.join( root, "1bos.mtz" )
    pdb_file = os.path.join( root, "monomer.pdb" )
    
    ( xray_data, symmetry ) = read_data( mtz_file, "FOBS_X", "SIGFOBS_X" )
    component = process_model.Component(
        sequence = "TPDCVTGKVEYTKYNDDDTFTVKVGDKELFTNRWNLQSLLLSAQITGMTVTIKTNACHNGGGFSEVIFR",
        )
    composition = process_model.Composition.FromComponents( components = [ component ] * 20 )
    mr_model = process_model.MRModel.Rms( pdb_file = pdb_file, value = 0.8 )
    mr_ensemble = process_model.MREnsemble(
        component = component,
        models = [ mr_model ],
        score = 0.45
        )
    problem = process_model.MRProblem(
        xray_data = xray_data,
        symmetry = symmetry,
        composition = composition
        )
    solver = get_mr_solver(
        selection_strategy = strategy.mr_set.DynamicGlobalTop(),
        engine = engine.Realtime(),
        analysers = [
            analyser.SolutionIdentify(),
            analyser.CombinatorialExtension(),
            analyser.SpeculativeCalculationAssess(),
            ]
        )
    loggers = get_loggers()
    blackboard = setup_solution_blackboard(
        startup_objects = [ problem, component, mr_ensemble, process_model.CPU() ]
        )
    
    for logger in loggers:
        blackboard.register( logger )
        
    solver.run( blackboard = blackboard )
    
    
def proteosome():
    
    root = "/home/gaborb/data/proteosome/"
    mtz_file = os.path.join( root, "scale.mtz" )
    pdb_alpha = os.path.join( root, "alpha.pdb" )
    pdb_beta = os.path.join( root, "beta.pdb" )
    
    ( xray_data, symmetry ) = read_data( mtz_file, "scale", "SIGscale" )
    component1 = process_model.Component(
        sequence = ( "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAACAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" )
        )
    component2 = process_model.Component(
        sequence = ( "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
            + "AAAAAAAAAAAAAAAAAA" )
        )
    composition = process_model.Composition.FromComponents(
        components = [ component1 ] * 7 + [ component2 ] * 7
        )
    mr_model_alpha = process_model.MRModel.Rms( pdb_file = pdb_alpha, value = 1.0 )
    mr_model_beta = process_model.MRModel.Rms( pdb_file = pdb_beta, value = 1.0 )
    mr_ensemble_alpha = process_model.MREnsemble(
        component = component1,
        models = [ mr_model_alpha ],
        score = 0.45
        )
    mr_ensemble_beta = process_model.MREnsemble(
        component = component2,
        models = [ mr_model_beta ],
        score = 0.55
        )
    problem = process_model.MRProblem(
        xray_data = xray_data,
        symmetry = symmetry,
        composition = composition
        )
    
    solver = get_mr_solver(
        job_setup_strategy = strategy.mr_set.GlobalTop(),
        additional_actors = [
            actors.SolutionIdentify(),
            actors.SearchThreadTerminate()
            ]
        )
    loggers = get_loggers()
    blackboard = setup_solution_blackboard(
        startup_objects = [
            problem,
            component1, component2,
            mr_ensemble_alpha, mr_ensemble_beta,
            ]
        )
    
    for logger in loggers:
        blackboard.register( logger )
        
    solver.run( blackboard = blackboard )
    
if __name__ == "__main__":
    lysozyme()