#!/usr/bin/python
from PIPRes.testing import *
from cStringIO import StringIO
from  garli_wrap import *
_LOG = cipresGetLogger('test.garli_wrap')

class GarliWrapTest(CipresTestingClient):
    serviceRegistryEntry = composeRegistryQuery(interface='TreeImprove', applicationName='GARLI')
    serviceType = CipresIDL_api1.TreeImprove
    def cipresTestSample(self):
        fn = '49Tx2543C.nex'
        p = getCipresTestFile(fn)
        if p is None:
            _LOG.warn('Could not find file %s in CIPRES test file directory.\nMake sure CIPRES_ROOT is set correctly and the file exists.' %s)
        firstSet =  readFilePathAsNexus(p)[0]
        mat, tree = firstSet.matrixList[0], firstSet.treeList[0]
        excmd = 'randseed = 10 ; availablememeory = 512  ; refinestart = 1 ; enforcetermconditions = 0 ; stopgen = 5000000 ; stoptime = 1 ; dontinferproportioninvariant = 1 ; nindivs = 4 ; holdover = 1 ; selectionintensity = 0.5 ; holdoverpenalty = 0 ; topoweight = 1.0 ; randnniweight = 1 ; randsprweight = 1 ; limsprweight = 1 ; limsprrange = 6 ; modweight = 0.05 ; gammashapemodel = 1000 ; brlenweight = 0.2 ; meanbrlenmuts = 5 ; gammashapebrlen = 1000 ; intervallength = 100 ; intervalstostore = 5 ; startoptprec = 0.5 ; minoptprec = 0.01 ; numberofprecreductions = 40 ;'
        tiObjRef = self.getTestedService()
        try:
            tiObjRef.setMatrix(mat)
            startTree = CipresTree('(12:.05,5:0.05,(6:0.06,(1:0.05,20:0.02):.04):.01)')
            tiObjRef.setTree(startTree)
            self.assertEqual(tiObjRef.execute(excmd), (True, ''))
            rTree = CipresTree(tiObjRef.improveTree(CORBA.Object._nil))
            expectedTree = CipresTree('(6:0.06552048,5:0.05023923,(12:0.04708402,(1:0.14642338,20:0.0704683):1e-08):0.03152402);')
            _LOG.info('testSample may fail because it relies on stoptime (the time should be short enough to prevent optimization, but fast machines may see an improvement in the tree)')
            self.assertEqual(rTree, expectedTree)
        finally:
            tiObjRef.remove()
    def cipresTestStoredFile(self):
        if 'QUICK_CIPRES_TESTS' in os.environ:
            _LOG.info('\nGarliWrapTest.testStoredFile skipped because QUICK_CIPRES_TESTS is in the env.')
            return
        m = CipresIDL_api1.DataMatrix('ACGT?', 4, 26, [[0], [1], [2], [3], [-1,0,1,2,3]],
            [   [0,1,2,3,4,0,1,2,3,4,1,1,3,1,1,2,3,2,0,1,2,3,0,1,2,3],
                [2,3,4,0,1,2,3,4,0,1,1,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3],
                [0,1,2,3,4,0,1,2,3,3,0,1,3,0,1,2,3,2,0,1,2,3,0,1,2,3],
                [2,3,4,0,1,2,3,4,1,1,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3],
            ], CipresIDL_api1.DNA_DATATYPE)
        t = CipresTree('(1,2,(3,4))')
        addScore(t, None)
        tiObjRef = self.getTestedService()
        try:
            tiObjRef.setMatrix(m)
            tiObjRef.setTree(t)
            tiObjRef.execute('randseed = 485')
            rTree = CipresTree(tiObjRef.improveTree(CORBA.Object._nil))
            ct = numberedLeafTree('(1,3,(2,4))')
            self.assertEqual(approxEqual(rTree.m_score.doubleScore, -77.9649, tol=.1), True)
            self.assertEqual(ct, rTree)
        finally:
            tiObjRef.remove()
    def testParseGarliBestTreeStream(self):
        self.assertEqual(exampleBestTreeContents, parseGarliBestTreeStream(exampleBestTreeFile))
    def testGarliModel(self):
        self.assertEqual(str(GarliModel()), 'r 1.0 5.0 1.0 1.0 5.0 b 0.25 0.25 0.25 0.25 a 0.500000 p 0.000000 ') 

absGarliPath = expandPath('$GARLI_EXE')
if os.path.exists(absGarliPath):
    try:
        GarliWrap.debugging = True
    except: 
        pass
    else:
        GarliWrapTest.serviceFactory = lambda s, x: GarliWrap(x, absGarliPath)
        
exampleBestTreeContents = (-21812.91, 'r 2.2083458 8.1652996 2.3565644 0.92223979 16.481014 b 0.34774389 0.21249714 0.14980668 0.2899523 a 0.55113683 p 0.33375217', '((((((((18:0.06835205,(19:0.05306490,20:0.09193022):0.06937089):0.01689346,((((21:0.08459729,22:0.00449459):0.02501980,24:0.01571060):0.00865423,23:0.02590949):0.01533227,(25:0.07082882,26:0.05789506):0.01418653):0.05899064):0.01836833,((((2:0.05713207,(3:0.00441135,4:0.00012159):0.02333281):0.00698649,((7:0.01643537,6:0.01483097):0.01451553,5:0.02231806):0.01848974):0.00789437,1:0.08157706):0.05002343,((8:0.00352179,9:0.00207403):0.05362261,((10:0.03115259,(((13:0.00438012,14:0.00552822):0.00969489,(15:0.02723586,(16:0.00000001,17:0.00116821):0.02253866):0.00233145):0.00612875,11:0.05653261):0.00238905):0.00612336,12:0.04125784):0.04149618):0.01434602):0.02868053):0.00788635,(((28:0.07427360,30:0.05257953):0.01538761,(29:0.03063870,31:0.07979972):0.01389653):0.03912552,27:0.10278336):0.02699129):0.08131591,((33:0.00195451,32:0.00685806):0.02327930,((36:0.01759142,(38:0.00763102,(39:0.00716960,37:0.01927415):0.00114027):0.00922435):0.00558553,(35:0.01629215,34:0.01003344):0.01507966):0.02020950):0.03180166):0.00622077,((45:0.04968721,(((53:0.05203892,51:0.07567231):0.01606596,(60:0.05805693,(61:0.00611512,52:0.00621863):0.02149677):0.00501085):0.00445728,(((58:0.02598885,57:0.01577610):0.00406806,59:0.02944709):0.00523670,(56:0.04990643,(55:0.00576922,54:0.00979596):0.01298527):0.01202457):0.00618748):0.00256316):0.00345134,((44:0.04143879,(50:0.05321411,62:0.02775947):0.00201038):0.00229726,((64:0.00989547,63:0.01438299):0.01621131,(49:0.00699223,(47:0.00527064,(46:0.00286219,48:0.00108350):0.00301750):0.00590560):0.02178807):0.00488662):0.00429467):0.02601383):0.01339155,40:0.02173790):0.00815648,(42:0.00192885,43:0.00964725):0.02100080,41:0.03279430)')
exampleBestTreeFile = StringIO("""#nexus

begin trees;
translate
  1  temporariaDMH84R1,
  2  boyliiMVZ148929,
  3  luteiventris_MT_MVZ191016,
  4  luteiventris_WA_MVZ225749,
  5  muscosaMVZ149006,
  6  auroraMVZ13957,
  7  cascadaeMVZ148946,
  8  sylvaticaMVZ137426,
  9  sylvaticaDMH84R43,
  10  septentrionalesDCC3588,
  11  grylioMVZ175945,
  12  virgatipesMVZ175944,
  13  okaloosae,
  14  clamitansJSF1118,
  15  heckscheriMVZ164908,
  16  catesbianaX12841,
  17  catesbianaDMH84R2,
  18  maculataKU195258,
  19  vibicariaMVZ11035,
  20  warszewitshiiJSF1127,
  21  palmipesVenAMNHA118801,
  22  palmipesEcuKU204425,
  23  bwanaQCAZ13964,
  24  Sp_1_ecuadorQCAZ13219,
  25  vaillantiKU195299,
  26  julianiTNHC60324,
  27  sierramadrensisKU195181,
  28  psilonotaKU195119,
  29  tarahumaraeKU194596,
  30  zweifeliJAC7514,
  31  pustulosaJAC10555,
  32  pipiensJSF1119,
  33  pipiensY10945,
  34  dunniJSF1017,
  35  montezumaeJAC8836,
  36  sp_2_mex_JSF1106,
  37  chiricahuensisJSF1063,
  38  chiricahuensisJSF1092,
  39  subaquavocalis,
  40  palustrisJSF1110,
  41  areolataJSF1111,
  42  sevosaUSC8236,
  43  capitoSLU003,
  44  spectabilisJAC8622,
  45  forreriJSF1065,
  46  tlalociJSF1083,
  47  berlandieriJSF1136,
  48  neovolcanicaJSF960,
  49  blairiJSF830,
  50  omiltemanaJAC7413,
  51  magnaocularisJSF1073,
  52  yavapaiensisJSF1085,
  53  sp_7_JaliscoJSF1000,
  54  macroglossaJAC10472,
  55  macroglossaJSF7933,
  56  taylori286,
  57  sp_4_Panama,
  58  sp_5_CostaRichDMH86_210,
  59  sp_6_CostaRicaDMH86_225,
  60  sp_8_PueblaJAC9467,
  61  oncaLVT3542,
  62  sp_3_MichoacanJSF955,
  63  sphenocephalaUSC7448,
  64  utriculariaJSF845;
tree best = [&U][-21812.91][ r 2.2083458 8.1652996 2.3565644 0.92223979 16.481014 b 0.34774389 0.21249714 0.14980668 0.2899523 a 0.55113683 p 0.33375217 ]((((((((18:0.06835205,(19:0.05306490,20:0.09193022):0.06937089):0.01689346,((((21:0.08459729,22:0.00449459):0.02501980,24:0.01571060):0.00865423,23:0.02590949):0.01533227,(25:0.07082882,26:0.05789506):0.01418653):0.05899064):0.01836833,((((2:0.05713207,(3:0.00441135,4:0.00012159):0.02333281):0.00698649,((7:0.01643537,6:0.01483097):0.01451553,5:0.02231806):0.01848974):0.00789437,1:0.08157706):0.05002343,((8:0.00352179,9:0.00207403):0.05362261,((10:0.03115259,(((13:0.00438012,14:0.00552822):0.00969489,(15:0.02723586,(16:0.00000001,17:0.00116821):0.02253866):0.00233145):0.00612875,11:0.05653261):0.00238905):0.00612336,12:0.04125784):0.04149618):0.01434602):0.02868053):0.00788635,(((28:0.07427360,30:0.05257953):0.01538761,(29:0.03063870,31:0.07979972):0.01389653):0.03912552,27:0.10278336):0.02699129):0.08131591,((33:0.00195451,32:0.00685806):0.02327930,((36:0.01759142,(38:0.00763102,(39:0.00716960,37:0.01927415):0.00114027):0.00922435):0.00558553,(35:0.01629215,34:0.01003344):0.01507966):0.02020950):0.03180166):0.00622077,((45:0.04968721,(((53:0.05203892,51:0.07567231):0.01606596,(60:0.05805693,(61:0.00611512,52:0.00621863):0.02149677):0.00501085):0.00445728,(((58:0.02598885,57:0.01577610):0.00406806,59:0.02944709):0.00523670,(56:0.04990643,(55:0.00576922,54:0.00979596):0.01298527):0.01202457):0.00618748):0.00256316):0.00345134,((44:0.04143879,(50:0.05321411,62:0.02775947):0.00201038):0.00229726,((64:0.00989547,63:0.01438299):0.01621131,(49:0.00699223,(47:0.00527064,(46:0.00286219,48:0.00108350):0.00301750):0.00590560):0.02178807):0.00488662):0.00429467):0.02601383):0.01339155,40:0.02173790):0.00815648,(42:0.00192885,43:0.00964725):0.02100080,41:0.03279430);
end;
begin paup;
clear;
gett file=ranaGarli.best.tre storebr;
lset userbr nst=6 rmat=(2.20834577 8.16529958 2.35656441 0.92223979 16.48101410) base=(0.34774389 0.21249714 0.14980668) rates=gamma shape= 0.55113683 pinv= 0.33375217;
end;
""")

if __name__ == '__main__':
    addCipresTests(GarliWrapTest)
    main()