#!/usr/bin/env python -u
'''
pDMET: Density Matrix Embedding theory for Periodic Systems
Copyright (C) 2018 Hung Q. Pham. All Rights Reserved.
A few functions in pDMET are modifed from QC-DMET Copyright (C) 2015 Sebastian Wouters

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
import scipy.optimize as optimize

'''DFT XC library'''
GGA_X = {
'GGA_X_GAM'                    : 32 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
'GGA_X_HCTH_A'                 : 34 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
'GGA_X_EV93'                   : 35 , # E. Engel and S. H. Vosko, Phys. Rev. B 47, 13164 (1993)
'GGA_X_BCGP'                   : 38 , # K. Burke, A. Cancio, T. Gould, and S. Pittalis, ArXiv e-prints (2014), arXiv:1409.4834 [cond-mat.mtrl-sci]
'GGA_X_LAMBDA_OC2_N'           : 40 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_B86_R'                  : 41 , # I. Hamada, Phys. Rev. B 89, 121103 (2014)
'GGA_X_LAMBDA_CH_N'            : 44 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_LAMBDA_LO_N'            : 45 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
'GGA_X_HJS_B88_V2'             : 46 , # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
'GGA_X_Q2D'                    : 48 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
'GGA_X_PBE_MOL'                : 49 , # J. M. del Campo, J. L. G\'azquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
'GGA_X_AK13'                   : 56 , # R. Armiento and S. Kummel, Phys. Rev. Lett. 111, 036402 (2013)
'GGA_X_LV_RPW86'               : 58 , # K. Berland and P. Hyldgaard, Phys. Rev. B 89, 035412 (2014)
'GGA_X_PBE_TCA'                : 59 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
'GGA_X_PBEINT'                 : 60 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
'GGA_X_VMT84_GE'               : 68 , # A. Vela, J. C. Pacheco-Kato, J. L. G\'azquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
'GGA_X_VMT84_PBE'              : 69 , # A. Vela, J. C. Pacheco-Kato, J. L. G\'azquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
'GGA_X_VMT_GE'                 : 70 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
'GGA_X_VMT_PBE'                : 71 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
'GGA_X_N12'                    : 82 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
'GGA_X_SSB_SW'                 : 90 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Comput. Methods Sci. Eng. 9, 69 (2009)
'GGA_X_SSB'                    : 91 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
'GGA_X_SSB_D'                  : 92 , # M. Swart, M. Sol\'a, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
'GGA_X_BPCCAC'                 : 98 , # E. Br\'emond, D. Pilard, I. Ciofini, H. Chermette, C. Adamo, and P. Cortona, Theor. Chem. Acc. 131, 1184 (2012)
'GGA_X_PBE'                    : 101, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_X_PBE_R'                  : 102, # Y. Zhang and W. Yang, Phys. Rev. Lett. 80, 890 (1998)
'GGA_X_B86'                    : 103, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
'GGA_X_HERMAN'                 : 104, # F. Herman, J. P. V. Dyke, and I. B. Ortenburger, Phys. Rev. Lett. 22, 807 (1969)
'GGA_X_B86_MGC'                : 105, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
'GGA_X_B88'                    : 106, # A. D. Becke, Phys. Rev. A 38, 3098 (1988)
'GGA_X_G96'                    : 107, # P. M. W. Gill, Mol. Phys. 89, 433 (1996)
'GGA_X_PW86'                   : 108, # J. P. Perdew and W. Yue, Phys. Rev. B 33, 8800 (1986)
'GGA_X_PW91'                   : 109, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_X_OPTX'                   : 110, # N. C. Handy and A. J. Cohen, Mol. Phys. 99, 403 (2001)
'GGA_X_DK87_R1'                : 111, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
'GGA_X_DK87_R2'                : 112, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
'GGA_X_LG93'                   : 113, # D. J. Lacks and R. G. Gordon, Phys. Rev. A 47, 4681 (1993)
'GGA_X_FT97_A'                 : 114, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
'GGA_X_FT97_B'                 : 115, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
'GGA_X_PBE_SOL'                : 116, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
'GGA_X_RPBE'                   : 117, # B. Hammer, L. B. Hansen, and J. K. Norskov, Phys. Rev. B 59, 7413 (1999)
'GGA_X_WC'                     : 118, # Z. Wu and R. E. Cohen, Phys. Rev. B 73, 235116 (2006)
'GGA_X_MPW91'                  : 119, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
'GGA_X_AM05'                   : 120, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_X_PBEA'                   : 121, # G. K. H. Madsen, Phys. Rev. B 75, 195108 (2007)
'GGA_X_MPBE'                   : 122, # C. Adamo and V. Barone, J. Chem. Phys. 116, 5933 (2002)
'GGA_X_XPBE'                   : 123, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
'GGA_X_2D_B86_MGC'             : 124, # S. Pittalis, E. Rasanen, J. G. Vilhena, and M. A. L. Marques, Phys. Rev. A 79, 012503 (2009)
'GGA_X_BAYESIAN'               : 125, # J. J. Mortensen, K. Kaasbjerg, S. L. Frederiksen, J. K. Norskov, J. P. Sethna, and K. W. Jacobsen, Phys. Rev. Lett. 95, 216401 (2005)
'GGA_X_PBE_JSJR'               : 126, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
'GGA_X_2D_B88'                 : 127, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_2D_B86'                 : 128, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_2D_PBE'                 : 129, # J. G. Vilhena and M. A. L. Marques, unpublished (2014)
'GGA_X_OPTB88_VDW'             : 139, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_PBEK1_VDW'              : 140, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_OPTPBE_VDW'             : 141, # J. Klime\v{s}, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
'GGA_X_RGE2'                   : 142, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
'GGA_X_RPW86'                  : 144, # E. D. Murray, K. Lee, and D. C. Langreth, J. Chem. Theory Comput. 5, 2754 (2009)
'GGA_X_KT1'                    : 145, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 119, 3015 (2003)
'GGA_X_MB88'                   : 149, # V. Tognetti and C. Adamo, J. Phys. Chem. A 113, 14415 (2009)
'GGA_X_SOGGA'                  : 150, # Y. Zhao and D. G. Truhlar, J. Chem. Phys. 128, 184109 (2008)
'GGA_X_SOGGA11'                : 151, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
'GGA_X_C09X'                   : 158, # V. R. Cooper, Phys. Rev. B 81, 161104 (2010)
'GGA_X_OL2'                    : 183, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
'GGA_X_APBE'                   : 184, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
'GGA_X_HTBS'                   : 191, # P. Haas, F. Tran, P. Blaha, and K. Schwarz, Phys. Rev. B 83, 205117 (2011)
'GGA_X_AIRY'                   : 192, # L. A. Constantin, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. B 80, 035125 (2009)
'GGA_X_LAG'                    : 193, # L. Vitos, B. Johansson, J. Koll\'ar, and H. L. Skriver, Phys. Rev. B 62, 10046 (2000)
'GGA_X_PBEFE'                  : 265, # R. Sarmiento-P\'erez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
'GGA_X_CAP'                    : 270, # J. Carmona-Esp\'indola, J. L. G\'azquez, A. Vela, and S. B. Trickey, J. Chem. Phys. 142, 054105 (2015), 10.1063/1.4906606
'GGA_X_EB88'                   : 271, # P. Elliott and K. Burke, Can. J. Chem. 87, 1485 (2009)
'GGA_X_BEEFVDW'                : 285, # J. Wellendorff, K. T. Lundgaard, A. M\o{gelh\o{j}, V. Petzold, D. D. Landis, J. K. N\o{rskov}, T. Bligaard, and K. W. Jacobsen, }Phys. Rev. B 85, 235149 (2012)
'GGA_X_PBETRANS'               : 291, # Eric Bremond, I. Ciofini, and C. Adamo, Mol. Phys. 114, 1059 (2016)
'GGA_X_CHACHIYO'               : 298, # T. {Chachiyo and H. {Chachiyo}, }ArXiv e-prints (2017), arXiv:1706.01343 [cond-mat.mtrl-sci]
'GGA_X_WPBEH'                  : 524, # J. Heyd, G. E. Scuseria, and M. Ernzerhof, J. Chem. Phys. 118, 8207 (2003)
'GGA_X_HJS_PBE'                : 525, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_HJS_PBE_SOL'            : 526, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_HJS_B97X'               : 528, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
'GGA_X_ITYH'                   : 529, # H. Iikura, T. Tsuneda, T. Yanai, and K. Hirao, J. Chem. Phys. 115, 3540 (2001)
'GGA_X_SFAT'                   : 530, # A. Savin and H.-J. Flad, Int. J. Quantum Chem. 56, 327 (1995)
'GGA_X_SG4'                    : 533, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
'GGA_X_GG99'                   : 535, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
'GGA_X_PBEPOW'                 : 539, # Eric Bremond, J. Chem. Phys. 145, 244102 (2016)
'GGA_X_KGG99'                  : 544, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
'GGA_X_B88M'                   : 570, # E. Proynov, H. Chermette, and D. R. Salahub, J. Chem. Phys. 113, 10013 (2000)
}

GGA_C = {
'GGA_C_GAM'                    : 33 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
'GGA_C_BCGP'                   : 39 , # K. Burke, A. Cancio, T. Gould, and S. Pittalis, ArXiv e-prints (2014), arXiv:1409.4834 [cond-mat.mtrl-sci]
'GGA_C_Q2D'                    : 47 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
'GGA_C_ZPBEINT'                : 61 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
'GGA_C_PBEINT'                 : 62 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
'GGA_C_ZPBESOL'                : 63 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
'GGA_C_N12_SX'                 : 79 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
'GGA_C_N12'                    : 80 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
'GGA_C_REGTPSS'                : 83 , # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, L. A. Constantin, and J. Sun, Phys. Rev. Lett. 103, 026403 (2009)
'GGA_C_OP_XALPHA'              : 84 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_G96'                 : 85 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_PBE'                 : 86 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_OP_B88'                 : 87 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_FT97'                   : 88 , # M. Filatov and W. Thiel, Int. J. Quantum Chem. 62, 603 (1997)
'GGA_C_HCTH_A'                 : 97 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
'GGA_C_REVTCA'                 : 99 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
'GGA_C_TCA'                    : 100, # V. Tognetti, P. Cortona, and C. Adamo, J. Chem. Phys. 128, 034101 (2008)
'GGA_C_PBE'                    : 130, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_C_LYP'                    : 131, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
'GGA_C_P86'                    : 132, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
'GGA_C_PBE_SOL'                : 133, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
'GGA_C_PW91'                   : 134, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_C_AM05'                   : 135, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_C_XPBE'                   : 136, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
'GGA_C_LM'                     : 137, # D. C. Langreth and M. J. Mehl, Phys. Rev. Lett. 47, 446 (1981)
'GGA_C_PBE_JRGX'               : 138, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
'GGA_C_RGE2'                   : 143, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
'GGA_C_WL'                     : 147, # L. C. Wilson and M. Levy, Phys. Rev. B 41, 12930 (1990)
'GGA_C_WI'                     : 148, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
'GGA_C_SOGGA11'                : 152, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
'GGA_C_WI0'                    : 153, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
'GGA_C_SOGGA11_X'              : 159, # R. Peverati and D. G. Truhlar, J. Chem. Phys. 135, 191102 (2011)
'GGA_C_APBE'                   : 186, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
'GGA_C_OPTC'                   : 200, # A. J. Cohen and N. C. Handy, Mol. Phys. 99, 607 (2001)
'GGA_C_PBELOC'                 : 246, # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 86, 035130 (2012)
'GGA_C_PBEFE'                  : 258, # R. Sarmiento-P\'erez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
'GGA_C_OP_PW91'                : 262, # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
'GGA_C_PBE_MOL'                : 272, # J. M. del Campo, J. L. G\'azquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
'GGA_C_BMK'                    : 280, # A. D. Boese and J. M. L. Martin, J. Chem. Phys. 121, 3405 (2004)
'GGA_C_TAU_HCTH'               : 281, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
'GGA_C_HYB_TAU_HCTH'           : 283, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
'GGA_C_SG4'                    : 534, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
'GGA_C_SCAN_E0'                : 553, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
'GGA_C_GAPC'                   : 555, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014), pMID: 26580528
'GGA_C_GAPLOC'                 : 556, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014), pMID: 26580528
'GGA_C_ZVPBEINT'               : 557, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
'GGA_C_ZVPBESOL'               : 558, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
'GGA_C_TM_LYP'                 : 559, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
'GGA_C_TM_PBE'                 : 560, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
'GGA_C_W94'                    : 561, # L. C. Wilson, Chemical Physics 181, 337 (1994)
'GGA_C_CS1'                    : 565, # N. C. Handy a
}

'''Smaller set of XC to test'''
GGA_X = {

'GGA_X_PBE'                    : 101, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_X_B88'                    : 106, # A. D. Becke, Phys. Rev. A 38, 3098 (1988)
'GGA_X_PW91'                   : 109, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
'GGA_X_WC'                     : 118, # Z. Wu and R. E. Cohen, Phys. Rev. B 73, 235116 (2006)
'GGA_X_AM05'                   : 120, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
}

GGA_C = {
'GGA_C_PBE'                    : 130, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
'GGA_C_LYP'                    : 131, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
'GGA_C_P86'                    : 132, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
'GGA_C_AM05'                   : 135, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
'GGA_C_SCAN_E0'                : 553, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
}

def get_init_uvec(xc_type='PBE0', dft_HF=None):
    '''Initialize uvec corresponding to each type of the DF-like cost function'''
    if dft_HF is not None:
        HF_percent = dft_HF
    else:
        HF_percent = 1.0

    if xc_type == 'PBE0':
        uvec = [HF_percent]  + [0.] * 2
    elif xc_type == 'RSH-PBE0':
        uvec = [1.] + [1.]  + [0.] * 2
    elif xc_type == 'B3LYP':
        uvec = [HF_percent]  + [0.] * 4
    elif xc_type == 'CAMB3LYP':
        uvec = [HF_percent , .0000001, 1. , 0., 0., 0., 0.]
    elif xc_type == 'manyGGA':
        uvec = [HF_percent]  + [0.] * (len(GGA_X) + len(GGA_C))

    return uvec

def get_bounds(xc_type='PBE0', constraint=1, dft_HF=None):
    '''Prodive the bounds to optimize the DF-like cost function'''

    if constraint == 1:
        if xc_type == 'PBE0':
            lower = [-1.0] + [.0] * 2
            upper = [ 1.0] + [1.0] * 2
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'RSH-PBE0':
            lower = [ -1.] * 2 + [.0] * 2
            upper = [ 1.0] * 4
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'B3LYP':
            lower = [-1.0] + [.0] * 4
            upper = [ 1.0] + [1.0] * 4
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'CAMB3LYP':
            bounds = optimize.Bounds([0.0,.0000001,.0,.0,.0,.0,.0], [1.,1.0,1.,1.,1.,1.,1.])
        elif xc_type == 'manyGGA':
            lower = [-1.0] + [.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] + [.1] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    elif constraint == 2:
        if xc_type == 'PBE0':
            lower = [-1.0] * 3
            upper = [ 1.0] * 3
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'RSH-PBE0':
            lower = [-1] * 4
            upper = [ 1.0]* 4
        elif xc_type == 'B3LYP':
            lower = [-1.0] * 5
            upper = [ 1.0] * 5
            bounds = optimize.Bounds(lower, upper)
        elif xc_type == 'CAMB3LYP':
            bounds = optimize.Bounds([0.0,.0000001,.0,.0,.0,.0,.0], [1.,1.0,1.,1.,1.,1.,1.])
        elif xc_type == 'manyGGA':
            lower = [-1.0] + [-1.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] + [ 1.0] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    elif constraint == 3:
        if xc_type == 'manyGGA':
            lower = [ 0.0] * (len(GGA_X) + len(GGA_C))
            upper = [ 1.0] * (len(GGA_X) + len(GGA_C))
            bounds = optimize.Bounds(lower, upper)
    return bounds

def get_OEH_kpts(local, umat, xc_type='PBE0', dft_HF=None):
    '''Construct the mf Hamiltonian'''

    if dft_HF is not None and xc_type is not 'RSH-PBE0':
        umat[0] = np.float64(dft_HF)

    if xc_type == 'PBE0':
        xc =  "{:.12f}*HF + {:.12f}*PBE, {:.12f}*PBE".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        if local._is_KROHF:
            OEH_kpts = 0.5 * (OEH_kpts[0] + OEH_kpts[1])
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'B3LYP':
        xc =  "{:.12f}*HF + {:.12f}*LDA + {:.12f}*B88, {:.12f}*LYP + {:.12f}*VWN".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        if local._is_KROHF:
            OEH_kpts = 0.5 * (OEH_kpts[0] + OEH_kpts[1])
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'RSH-PBE0':
        # TODO: right now the range-seperation parameter is fixed, it is super expensive to update it every cycle
        # because the vk long-range needed to be generated
        umat = [local.xc_omega] + np.asarray(umat).tolist()
        xc =  "{1:.12f}*SR_HF({0:.2f})+ {2:.12f}*LR_HF({0:.2f}) + {3:.12f}*PBE, {4:.12f}*PBE".format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - 0.5 * (umat[1]*local.vksr + umat[2]*local.vklr)
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'CAMB3LYP':
        # TODO: need to debug
        xc =  '{:.12f}*SR_HF({:.12f}) + {:.12f}*LR_HF({:.12f}) + {:.12f}*ITYH + {:.12f}*B88, {:.12f}*VWN5 + {:.12f}*LYP'.format(umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    elif xc_type == 'manyGGA':
        x = '{:.12f}*HF + {:.12f}*' + '+ {:.12f}*'.join(GGA_X.keys())
        c = ', {:.12f}*' + '+ {:.12f}*'.join(GGA_C.keys())
        xc = (x + c).format(*umat)
        n, exc, vxc = local.kks._numint.nr_rks(local.cell, local.kks.grids, xc, local.dm_kpts, 0, local.kpts, None)
        veff = vxc + local.vj - umat[0] * 0.5 * local.vk
        OEH_kpts = local.h_core + veff
        OEH_kpts = local.ao_2_loc(OEH_kpts, local.ao2lo)

    return OEH_kpts


from keras import backend as K
import os
import time
import h5py
import sys
from tagger_net import MusicTaggerCRNN
from keras.optimizers import SGD
import numpy as np
from keras.utils import np_utils
from math import floor
from music_tagger_cnn import MusicTaggerCNN
from utils import save_data, load_dataset, save_dataset, sort_result, predict_label, extract_melgrams
import matplotlib.pyplot as plt

# Parameters to set
TEST = 1

LOAD_MODEL = 0
LOAD_WEIGHTS = 1
MULTIFRAMES = 1
time_elapsed = 0

# GTZAN Dataset Tags
tags = ['blues', 'classical', 'country', 'disco', 'hiphop', 'jazz', 'metal', 'pop', 'reggae', 'rock']
tags = np.array(tags)

# Paths to set
model_name = "example_model"
model_path = "models_trained/" + model_name + "/"
weights_path = "models_trained/" + model_name + "/weights/"

test_songs_list = 'list_example.txt'

# Initialize model
model = MusicTaggerCRNN(weights=None, input_tensor=(1, 96, 1366))

model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

if LOAD_WEIGHTS:
    model.load_weights(weights_path+'crnn_net_gru_adam_ours_epoch_40.h5')

#model.summary()

X_test, num_frames_test= extract_melgrams(test_songs_list, MULTIFRAMES, process_all_song=False, num_songs_genre='')

num_frames_test = np.array(num_frames_test)

t0 = time.time()

print '\n--------- Predicting ---------','\n'

results = np.zeros((X_test.shape[0], tags.shape[0]))
predicted_labels_mean = np.zeros((num_frames_test.shape[0], 1))
predicted_labels_frames = np.zeros((X_test.shape[0], 1))

song_paths = open(test_songs_list, 'r').read().splitlines()

previous_numFrames = 0
n=0
for i in range(0, num_frames_test.shape[0]):
    print 'Song number' +str(i)+ ': ' + song_paths[i]

    num_frames=num_frames_test[i]
    print 'Num_frames of 30s: ', str(num_frames),'\n'

    results[previous_numFrames:previous_numFrames+num_frames] = model.predict(
        X_test[previous_numFrames:previous_numFrames+num_frames, :, :, :])

    s_counter = 0
    for j in range(previous_numFrames, previous_numFrames+num_frames):
        #normalize the results
        total = results[j,:].sum()
        results[j,:]=results[j,:]/total
        print 'Percentage of genre prediction for seconds '+ str(20+s_counter*30) + ' to ' \
            + str(20+(s_counter+1)*30) + ': '
        sort_result(tags, results[j,:].tolist())

        predicted_label_frames=predict_label(results[j,:])
        predicted_labels_frames[n]=predicted_label_frames
        s_counter += 1
        n+=1

    print '\n', 'Mean genre of the song: '
    results_song = results[previous_numFrames:previous_numFrames+num_frames]

    mean=results_song.mean(0)
    sort_result(tags, mean.tolist())

    predicted_label_mean=predict_label(mean)

    predicted_labels_mean[i]=predicted_label_mean
    print '\n','The predicted music genre for the song is', str(tags[predicted_label_mean]),'!\n'

    previous_numFrames = previous_numFrames+num_frames

    print '************************************************************************************************'

colors = ['b','g','c','r','m','k','y','#ff1122','#5511ff','#44ff22']
fig, ax = plt.subplots()
index = np.arange(tags.shape[0])
opacity = 1
bar_width = 0.2
print mean
#for g in range(0, tags.shape[0]):
plt.bar(left=index, height=mean, width=bar_width, alpha=opacity, color=colors)

plt.xlabel('Genres')
plt.ylabel('Percentage')
plt.title('Scores by genre')
plt.xticks(index + bar_width / 2, tags)
plt.tight_layout()
fig.autofmt_xdate()
plt.savefig('genres_prediction.png')


# Matthew Smith, 12/5/19
# Shuo Gu, 02/18/20
# Shoichet Lab: Torsion Strain Rotation Project
# This Python script will contain all of the functions for using the Torsion
# Library method for estimating torsion strain energy. To use these functions,
# source this script in either another Python script or a Python commandline
# session as:
#   exec(open("TL_Functions.py").read())
# when in the same directory.

# We will use TL_2.1_VERSION_6.xml as the Torsion Library.

# Explanation of how this works from TL_Lookup_Test.py:
# Since the Torsion Library contains very specific SMARTS patterns to define
# the torsion patterns (also called torsion rules in the acutal library) and
# there are only 514 torsion patterns in TL_2.1, it is best to just search
# over all of the torsion rules in the library, instead of trying to find
# and classify every substructure in the molecule. This was Ying's suggestion.
# For each torsion rule, we will use the GetSubstructMatches function in
# the package rdkit.Chem to find the indeces of the atoms that match the
# torsion rule. If there are multiple such substructures, then this function
# should return multiple sets of indeces. We will then find the information
# we need (dihedral angle, method, estimate, lower bound, upper bound, and
# flag). Since the same torsion pattern in the molecule can have multiple
# torsion rules in the library applying to it, we need to choose the intended
# torsion rule. The authors of the original Torsion Library ordered the
# torsion rules within each subclass in decreasing specificity, so we will
# keep track of the number of the torsion rule for each matching substructure.
# We will then pick the set of information with the lowest counting index,
# assuming that this is the most specific. Since the first hierarchy class in
# the Torsion Library is the general class, we will need to loop over all of
# the specific classes first. This process is not the most efficient, since
# we are looking up and calculating information that we immediately throw away,
# but we can optimize this procedure later.

# Once we have one set of information for each set of indeces defining the
# bonds, we need to condense these by torsion pattern. We must do this because,
# for a given bond, each atom can be bonded to 2 or 3 other atoms in a torsion
# pattern. This gives between 2*2=4 and 3*3=9 possible lists of 4 atom indeces
# defining the torsion pattern, not including reversing the order of the
# indeces! Therefore we need to find all of the information sets (angle,
# estimate, etc.) with the same atom indeces defining the bond (listed
# forwards and backwards) and decide which one we will keep. Since we are
# using these energy estimates for identifying unstable conformations in
# DOCK-ing campaigns, being conservstive with our hit-picking means being
# liberal in our energy assignments. We thus keep the set of information with
# the highest energy estimate, prioritizing the estimates from specific
# hierarchy classes over ones from the general hierarchy class.

# Shuo Gu, 01/10/20, fixed a bug of input reading
# Shuo Gu, 01/12/20, Introduced interpolation between energy bins
# Shuo Gu, 02/18/20, Changed the initial energy 0 to 100 for the approximate method

import xml.etree.ElementTree as ET #For reading XML files
from rdkit import Chem #The module with the RDKit functions we will need
import os #For the function to read in the molecules from the .mol2 file
import numpy as np #For arrays and NumPy function
from math import sqrt, atan2, pi #For calculating dihedral angles
from math import ceil #Ceiling function

# Import the XML file, using this as a guide:
# https://docs.python.org/3/library/xml.etree.elementtree.html
tree = ET.parse("TL_2.1_VERSION_6.xml")
root = tree.getroot()

# We will turn the procedure for estimating torsion strain energy in
# TL_Lookup_Test.py into a function that we can call on every molecule
# in the list, which will return an object of the following class:
class TP_list(object):
    def __init__(self, indeces, angles, smarts, hc, methods, E,
    CI_l, CI_u, flags):
        self.indeces = indeces #List of lists of indeces
        self.angles = angles #List of angles
        self.smarts = smarts #List of SMARTS strings
        self.hc = hc #List of strings for hierarchy class
        self.methods = methods #List of strings of energy-estimation methods
        self.E = E #List of energy estimates
        self.CI_l = CI_l #List of lower bounds for 95% CI of energy estimates
        self.CI_u = CI_u #List of upper bounds for 95% CI of energy estimates
        self.flags = flags
        # List of Booleans flagging whether an angle whose "method" is
        # "approximate" is not observed
        self.TP_indeces = [j for j in range(len(indeces))]
        # Indeces of the torsion patterns, not the atoms!

    # Getters and setters:
    def get_indeces(self):
        return(self.indeces)
    def set_indeces(self, inds):
        self.indeces = inds
    def get_angles(self):
        return(self.angles)
    def set_angles(self, angs):
        self.angles = angs
    def get_smarts(self):
        return(self.smarts)
    def set_smarts(self, sms):
        self.smarts = sms
    def get_hc(self):
        return(self.hc)
    def set_hc(self, hcs):
        self.hc = hcs
    def get_methods(self):
        return(self.methods)
    def set_methods(self, meths):
        self.methods = meths
    def get_E(self):
        return(self.E)
    def set_E(self, Es):
        self.E = Es
    def get_CI_l(self):
        return(self.CI_l)
    def set_CI_l(self, ls):
        self.CI_l = ls
    def get_CI_u(self):
        return(self.CI_u)
    def set_CI_u(self, us):
        self.CI_u = us
    def get_flags(self):
        return(self.flags)
    def set_flags(self, fs):
        self.flags = fs
    def get_TP_indeces(self): #Don't need a setter for this
        return(self.TP_indeces)

    # A method to return the information for a subset of the torsion patterns
    def get_TPs(self, inds = None):
        # The parameter inds is a list of indeces for the torsion patterns,
        # not the atoms! We default to using all of the torsion patterns
        if inds == None:
            inds = [j for j in range(len(self.indeces))]
        # Create a list of torsion pattern info to be returned:
        tps = [] #Initialize
        for j in inds:
            tps.append([self.TP_indeces[j], self.E[j],
            self.CI_l[j], self.CI_u[j], self.indeces[j], self.angles[j],
            self.smarts[j], self.hc[j], self.methods[j], self.flags[j]])
        return(tps)

    # A method to find the sum of the energy estimates, with the 95% CI
    def sum(self, cutoff = None):
        # If at least one of the torsion patterns is flagged, then we
        # do not calculate the estimate or CI. We return -1 * (number flagged)
        flagged = sum(self.flags) #Treats True = 1 and False = 0
        if flagged > 0:
            return([-1*flagged, 0, 0]) #We do not try to calculate the CI
        else:
            if cutoff == None:
                cutoff == 0 #The default cutoff is 0, or using every angle
            ret = [0] * 3 #Initialize the returned array
            for i in range(len(self.E)):
                if self.E[i] >= cutoff:
                    # If the energy estimate is larger than the cutoff
                    ret[0] += self.E[i]
                    ret[1] += self.CI_l[i]
                    ret[2] += self.CI_u[i]
            return(ret)

# We need a function to read in .mol2 files, which RDKit does not supply to us
# We will use a function from
# https://chem-workflows.com/articles/2019/07/18/building-a-multi-molecule-mol2-reader-for-rdkit/
# Shuo has modified this function significantly to get it to work with the
# .mol2 files that this lab uses, which contain many lines of comments before
# each line "@<TRIPOS>MOLECULE"
def Mol2MolSupplier (file = None):
    names = [] #Make a list to hold the molecule names
    mols = {} #Make a dictionary
    with open(file, 'r') as f:
        fileend = os.fstat(f.fileno()).st_size
        count = 0
        line = f.readline()

        while not f.tell() == fileend:
            if line.startswith("#") or line == '\n':
                line = f.readline()
            if line.startswith("@<TRIPOS>MOLECULE"):
                count += 1
                mol = []
                mol.append(line)
                line = f.readline()
                if line != "\n" and line.split()[0].strip() not in names:
                    name = line.split()[0].strip()
                    print(name)
                else:
                    name = "mol2Number" + str(count)
                    print(name)

                while not line.startswith("@<TRIPOS>MOLECULE"):
                    mol.append(line)
                    line = f.readline()

                    if f.tell() == fileend:
                        mol.append(line)
                        break
                block = ",".join(mol).replace(',','')
                m = Chem.rdmolfiles.MolFromMol2Block(block, sanitize=False, removeHs = False)
                names.append(name)
                mols[name] = m
    return(names, mols)

# Here is an updated version to use with the output "file" string buffer object
# created by the db2_file_like function in the db2_to_mol2.py script
def db2MolSupplier(file):
    names = [] #Make a list to hold the molecule names
    mols = {} #Make a dictionary
    with file as f: #file is already opened as a string buffer
        bufferend = len(f.getvalue())
        count = 0
        line = f.readline()
        while not f.tell() == bufferend:
            if line.startswith("#") or line == '\n':
                line = f.readline()
            if line.startswith("@<TRIPOS>MOLECULE"):
                count += 1
                mol = []
                mol.append(line)
                line = f.readline()

                name = "db2Number" + str(count)

                while not line.startswith("@<TRIPOS>MOLECULE"):
                    mol.append(line)
                    line = f.readline()

                    if f.tell() == bufferend:
                        mol.append(line)
                        break
                block = ",".join(mol).replace(',','')
                m = Chem.rdmolfiles.MolFromMol2Block(block, sanitize=False, removeHs = False)
                names.append(name)
                mols[name] = m
    return(names, mols)

# The first function we will need to make TP_list objects is the dihedral angle
# calculator. I made this function (and its helper function) for my iPQB
# summer project of writing a script that creates Ramachandran plots
# (Make_Rama.py). I have changed some of the comments explaining the math
# First a helper function for creating unit vectors
def unit(a):
    # The argument should be a NumPy array with 1 axis
    return(a / sqrt(np.dot(a,a))) #Scales a by its norm

def dihedral(a_1, a_2, a_3, a_4):
    # The arguments should all be NumPy arrays with 1 axis and length 3
    # These atoms should be in order, with a_2 and a_3 defining the bond
    # of interest

    # The 3 displacement vectors:
    b_1 = a_2 - a_1
    b_2 = a_3 - a_2
    b_3 = a_4 - a_3

    n_1 = unit(np.cross(b_1, b_2))
    n_2 = unit(np.cross(b_2, b_3))

    # Imagine the first atom (a_1) is above the middle bond (from a_2 to a_3),
    # so that b_1 points downward. Then n_1 points out of the page

    m = unit(np.cross(n_1, b_2))
    # I moved the normalization to be after the cross product. Moving
    # the normalization should not change the end result because the cross
    # product commutes with scalar multiplication and ||n_1|| = 1

    # Looking down b_2, we can consider n_1 to be the x-axis and
    # m to be the y-axis. Then the dihedral angle is the angle that
    # n_2 makes with the x-axis when projected into this plane. Since dihedral
    # angles are measured going clockwise, we need to negate the angle
    # that we get back from atan
    x = np.dot(n_1, n_2) #Project n_2 onto n_1
    y = np.dot(m, n_2) #Project n_2 onto m
    return(-atan2(y,x) * 180 / pi) #Return the angle in degrees

# The next function we will need will calculate angular differences,
# in degrees. This function will calculate theta_1 - theta_2. See my notes
# from 9/19/19
def ang_diff(theta_1, theta_2):
    # (-180,180] -> [0, 360)
    if theta_1 < 0:
        theta_1 += 360
    if theta_2 < 0:
        theta_2 += 360
    del_theta = (theta_1 - theta_2) % 360 #Angular difference
    # [0, 360) -> (-180, 180]
    if del_theta > 180:
        del_theta -= 360
    return(del_theta)
    # Test this works: ang_diff(0, 179), ang_diff(0, -179),
    # and ang_diff(-179, 179)

# This function will allow us to do the matching for each torsion rule
def tp_match(tp, hc, j, mol, pos, bi):
    # tp is a torsion pattern, hc is the type of hierarchyClass ("general" or
    # "specific", and j is the current value for i
    # This function turned the global variables mol, positions, and bond_info
    # in TL_Lookup_Test.py into parameters mol, pos and bi, respectively
    smarts = tp.get("smarts")
    # Create the histograms for energy estimates and bounds of confidence
    # intervals, if available
    hist_E = [] #Initialize for energy estimates
    hist_l = [] #Initialize for lower bounds of CIs
    hist_u = [] #Initialize for upper bounds of CIs
    if tp.get("method") == "exact":
        for bin in tp.find("histogram_converted").findall("bin"):
            hist_E.append(float(bin.get("energy")))
            hist_l.append(float(bin.get("lower")))
            hist_u.append(float(bin.get("upper")))
    matches = mol.GetSubstructMatches(Chem.MolFromSmarts(smarts))
    # A list of lists
    for match in matches: #For each match
        # Some of the SMARTS for the torion patterns actually have 5 atoms.
        # We need to ingore these
        if len(match) > 4:
            continue #Go to the next match

        # First get the atom locations
        # Based on https://github.com/rdkit/rdkit/issues/1982
        r_1 = np.array(pos[match[0]])
        r_2 = np.array(pos[match[1]])
        r_3 = np.array(pos[match[2]])
        r_4 = np.array(pos[match[3]])

        theta = dihedral(r_1, r_2, r_3, r_4) #Dihedral angle

        # Changed next line from "TP.get" to "tp.get"
        if tp.get("method") == "exact": #If using the exact method
            # First figure out what bin we are in in the histogram.
            # We define the bins by the right endpoints, like we did
            # when we made the plots of the energy profiles
            bin_num = ceil(theta / 10) + 17
            # Starting with the -170 deg bin
            energy = (hist_E[bin_num]-hist_E[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_E[bin_num]
            lower = (hist_l[bin_num]-hist_l[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_l[bin_num]
            upper = (hist_u[bin_num]-hist_u[(bin_num+35)%36])/10.0*(theta-(bin_num-17)*10)+hist_u[bin_num]

            bi.append(
            [
            list(match), #Convert tuple to list
            theta,
            smarts,
            hc, #"general" or "specific"
            "exact",
            energy,
            lower,
            upper,
            False, #This only could apply for the approximate method
            j #We will take this out when we create the final object
            ]
            )

        else: #If using the approximate method
            not_observed = True
            # Initialize the flag for not observing that angle
            energy = 100
            # Initialize the energy, in case we cannot estimate it
            # Loop over all the possible angle peaks
            for angle in tp.find("angleList").findall("angle"):
                theta_0 = float(angle.get("theta_0")) #Peak location
                delta = ang_diff(theta, theta_0) #Angular displacement
                if abs(delta) <= float(angle.get("tolerance2")):
                    # If within the tolerance range for that peak
                    beta_1 = float(angle.get("beta_1"))
                    beta_2 = float(angle.get("beta_2"))
                    # The coefficients for the regression
                    energy = beta_1*(delta ** 2) + beta_2*(delta ** 4)
                    # Using the "not-as-small angle approximation"
                    not_observed = False
                    break
                    # Break the for loop to avoid problems if the
                    # observed angle sits at the border between two
                    # peaks
            bi.append(
            [
            list(match), #Convert tuple to list
            theta,
            smarts,
            hc, #"general" or "specific"
            "approximate",
            energy,
            energy, #Lower bound for CI, which we cannot find for approx. method
            energy, #Upper bound for CI, which we cannot find for approx. method
            not_observed,
            j #We will take this out when we create the final object
            ]
            )

# This most general function will automate what we did in TL_Lookup_Test.py
def TL_lookup(mol): #mol is read in from the .mol2 file
    positions = mol.GetConformer().GetPositions()
    # List of lists of atom coordinates. Luckily RDKit starts indexing at 0
    bond_info = []
    # Initialize an empty list that will hold the information for each bond
    i = 0 #Initialize count of torsion rules

    # Loop over all of the specific hierarchy classes
    for HC in root.findall("hierarchyClass"):
        if HC.get("name") != "GG": #Not the general class
            for TP in HC.iter("torsionRule"): #Loop over each torsion rule
                tp_match(TP, "specific", i, mol, positions, bond_info)
                i += 1 #Increase the count for the torsion rule

    # Now for the general method:
    for TP in root.find("hierarchyClass[@name='GG']").iter("torsionRule"):
        tp_match(TP, "general", i, mol, positions, bond_info)
        i += 1 #Increase the count for the torsion rule

    # Now that we have all of the torsion patterns, we need to be able to find
    # duplicates. The first such way is if the entire pattern is reversed. We
    # can fix this problem by making sure that all of the lists of indeces have
    # the second index (the first in the bond of interest) lower than the third
    # index (the second in the bond of interest)
    for  bond in bond_info: #Loop over every bond
        if bond[0][1] > bond[0][2]:
            bond[0].reverse() #Reverse this list
            bond.append(True) #Mark that we reversed this bond's indeces
        else:
            bond.append(False) #Mark that we did not reverse this bond's
            # indeces. We will remove this marking later

    # Next we condense the bond_info by the lists of 4 atoms defining the bonds.
    # We will pick the entry of bond_info that has the lowest value for i for
    # each match, since the torsion rules in the Torsion Library are arranged
    # (within each hierarchy class or hierarchy subclass) in decreasing
    # specificity, and we loop over all of the specific hierarchy classes
    # before the general one
    bond_info_red = [bond_info[0]] #Initialize a list for the reduced bond info
    # This reduced list needs at least one element for checking subelements
    for j in range(1, len(bond_info)):
        # Skip the first bond, which is already in the reduced list
        atom_0 = bond_info[j][0][0] #First atom index
        atom_1 = bond_info[j][0][1] #Second atom index
        atom_2 = bond_info[j][0][2] #Third bond index
        atom_3 = bond_info[j][0][3] #Fourth bond index

        unmatched = True #Initialize not finding a match
        for k in range(len(bond_info_red)):
            # Check against everything in the growing reduced list
            if bond_info_red[k][0][0] == atom_0 \
            and bond_info_red[k][0][1] == atom_1 \
            and bond_info_red[k][0][2] == atom_2 \
            and bond_info_red[k][0][3] == atom_3:
                # If there is a match in ALL of the atom indeces
                unmatched = False
                if bond_info[j][9] < bond_info_red[k][9]:
                    # The ninth index gives the value for i
                    # If the new bond has a lower value, then we use it
                    # to replace the current one
                    bond_info_red[k] = bond_info[j]
                    break
                    # No need to continue looking for matches, since there
                    # should be no more than 1

        if unmatched: #If no match
            bond_info_red.append(bond_info[j]) #Append the current bond

    # Now we condense the bond_info_red by the 2 atoms actually defining the
    # bond. We will pick the entry of bond_info_red for each match that has
    # the highest energy estimate, prioritizing torsion patterns from
    # "specific" hierarchy classes over ones from the "general" hierarchy class
    b_i_r = [bond_info_red[0]] #Initialize a list for the further reduced bond info
    # I used the name b_i_r over bond_info_red_2
    for j in range(1, len(bond_info_red)):
        # Skip the first bond, which is already in the reduced list
        atom_1 = bond_info_red[j][0][1] #First atom index of the bond
        atom_2 = bond_info_red[j][0][2] #Second atom index of the bond

        unmatched = True #Initialize not finding a match
        for k in range(len(b_i_r)):
            # Check against everything in the growing list
            if b_i_r[k][0][1] == atom_1 and b_i_r[k][0][2] == atom_2:
                # If there is a match in the two atom indeces defining the bond
                unmatched = False
                if bond_info_red[j][3][0] > b_i_r[k][3][0] \
                or (
                bond_info_red[j][5] > b_i_r[k][5] \
                and
                bond_info_red[j][3][0] == b_i_r[k][3][0]
                ):
                    # The third index gives "general" or "specific", and we
                    # access the first char in that string. We use:
                    #   's' > 's' gives False
                    #   'g' > 'g' gives False
                    #   'g' > 's' gives False
                    #   's' > 'g' gives True
                    # to prioritize using the specific classe over the general one.
                    # The fifth index gives the energy estimate
                    b_i_r[k] = bond_info_red[j] #Replace the current bond info
                    break
                    # No need to continue looking for matches, since there
                    # should be no more than 1

        if unmatched: #If no match
            b_i_r.append(bond_info_red[j]) #Append the current bond

    # Now that we have only the bond information that we want, we need to
    # un-do any times we reversed the atom indeces
    for  bond in b_i_r:
        if bond[10]:
            # The tenth index gives whether or not we reversed the indeces
            bond[0].reverse() #Reverse this list

    # Now that we have all the information we need, we can return the
    # object that contains it
    return(
    TP_list([bond[0] for bond in b_i_r], [bond[1] for bond in b_i_r],
    [bond[2] for bond in b_i_r], [bond[3] for bond in b_i_r],
    [bond[4] for bond in b_i_r], [bond[5] for bond in b_i_r],
    [bond[6] for bond in b_i_r], [bond[7] for bond in b_i_r],
    [bond[8] for bond in b_i_r]) #Using Python list comprehension
    )


import argparse
import math
import time
import os
import shutil

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
import numpy as np

from dataloader import listflowfile as lt
from dataloader import SceneFlowLoader_demo as DA
from models.StereoNet import StereoNet

default_param = {
    'scale num':
    4,
    'max disparity':
    192,
    'batch size':
    2,
    'load model':
    'model_save_flyingthings3d_flyingthins3dTEST_b2_diff_scale_4/checkpoint_14.pth',
    'data path':
    '/home/caitao/Dataset/sceneflow/',
}
description_experiment = {
    'train data':
    'flyingthins3d',
    'test data':
    'flyingthings3d-TEST',
    'model':
    'StereoNet using difference to generate cost volume with scale is equal to four',
}
__normalize = {'mean': [0.0, 0.0, 0.0], 'std': [1.0, 1.0, 1.0]}

parser = argparse.ArgumentParser(description=description_experiment['model'])

parser.add_argument(
    '--maxdisp',
    type=int,
    default=default_param['max disparity'],
    help='maxium disparity')
# scale of low-solution image feature to input image
parser.add_argument(
    '--scalenum',
    type=int,
    default=default_param['scale num'],
    help='scale num')
parser.add_argument(
    '--datapath', default=default_param['data path'], help='datapath')
parser.add_argument(
    '--batchsize',
    type=int,
    default=default_param['batch size'],
    help='batch size')
parser.add_argument(
    '--loadmodel', default=default_param['load model'], help='load model')
parser.add_argument(
    '--no-cuda',
    action='store_true',
    default=False,
    help='enables CUDA training')
parser.add_argument(
    '--seed',
    type=int,
    default=1,
    metavar='S',
    help='random seed (default: 1)')

args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

_, _, _, test_left_img, test_right_img, test_left_disp = lt.dataloader(
    args.datapath)

TestImgLoader = torch.utils.data.DataLoader(
    DA.myImageFloder(
        test_left_img,
        test_right_img,
        test_left_disp,
        False,
        normalize=__normalize),
    batch_size=args.batchsize,
    shuffle=False,
    num_workers=min(4, args.batchsize),
    drop_last=False)

model = StereoNet(args.scalenum, args.scalenum, args.maxdisp)

if args.cuda:
    model = nn.DataParallel(model)
    model.cuda()

def save_prediction_for_demo(prediction_list, save_dir, left_path, disp_path,
                             batch_idx):
    save_dir = os.path.join(save_dir, batch_idx)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    else:
        print(save_dir, ' is existed!')
        import sys
        sys.exit(-1)

    first = os.path.join(save_dir, '1')
    if not os.path.exists(first):
        os.mkdir(first)
    second = os.path.join(save_dir, '2')
    if not os.path.exists(second):
        os.mkdir(second)

    shutil.copyfile(left_path[0],
                    os.path.join(first, os.path.basename(left_path[0])))
    shutil.copyfile(left_path[1],
                    os.path.join(second, os.path.basename(left_path[1])))

    first_disp = torch.from_numpy(np.load(disp_path[0]))
    tmp = torch.where(first_disp.data < 0, torch.zeros_like(first_disp.data),
                      first_disp.data)
    tmp = torch.where(tmp > 255, torch.full_like(tmp, 255), tmp)
    tmp /= 255
    first_disp = transforms.ToPILImage()(torch.unsqueeze(tmp, dim=0).cpu())
    first_disp.save(os.path.join(first, f'gt.png'))

    second_disp = torch.from_numpy(np.load(disp_path[1]))
    tmp = torch.where(second_disp.data < 0, torch.zeros_like(second_disp.data),
                      second_disp.data)
    tmp = torch.where(tmp > 255, torch.full_like(tmp, 255), tmp)
    tmp /= 255

    second_disp = transforms.ToPILImage()(torch.unsqueeze(tmp, dim=0).cpu())
    second_disp.save(os.path.join(second, f'gt.png'))

    length = len(prediction_list)
    for k, pred in enumerate(prediction_list):

        tmp = torch.where(pred.data < 0, torch.zeros_like(pred.data),
                          pred.data)
        tmp = torch.where(tmp > 255, torch.full_like(tmp, 255), tmp)

        tmp = tmp / 255

        prediction_img = transforms.ToPILImage()(torch.unsqueeze(
            tmp[0], dim=0).cpu())
        prediction_img.save(
            os.path.join(first, f'prediction_1_{math.pow(2, length-k-1)}.png'))

        prediction_img = transforms.ToPILImage()(torch.unsqueeze(
            tmp[1], dim=0).cpu())
        prediction_img.save(
            os.path.join(second,
                         f'prediction_1_{math.pow(2, length-k-1)}.png'))

def test(imgL, imgR, disp_true, save_dir, left_path, disp_path, batch_idx):
    with torch.no_grad():
        # model.eval()
        imgL = torch.FloatTensor(imgL)
        imgR = torch.FloatTensor(imgR)
        if args.cuda:
            imgL, imgR, disp_true = imgL.cuda(), imgR.cuda(), disp_true.cuda()

        # ---------
        mask = (disp_true < args.maxdisp) & (disp_true >= 0)
        # ----

        prediction_list = model(imgL, imgR)

        length = len(prediction_list)
        for i in range(length):
            if prediction_list[i].size()[-2:] != mask.size()[-2:]:
                assert i != length - 1
                prediction_list[i] *= (
                    mask.size()[-1] / prediction_list[i].size()[-1])
                prediction_list[i] = torch.squeeze(
                    F.interpolate(
                        torch.unsqueeze(prediction_list[i], dim=1),
                        size=mask.size()[-2:],
                        mode='bilinear',
                        align_corners=False),
                    dim=1)

        save_prediction_for_demo(prediction_list, save_dir, left_path,
                                 disp_path, str(batch_idx))
        output = prediction_list[-1]

        if len(disp_true[mask]) == 0:
            return 0
        else:
            # end-point-error
            return (torch.mean(
                torch.abs(output[mask] - disp_true[mask]))).item()

def test_model(model_path, save_dir, batch_count):

    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict['state_dict'])
    # ------------- TEST -----------------------------------------------------
    total_test_loss = 0
    for batch_idx, (imgL, imgR, disp_L, left_path,
                    disp_path) in enumerate(TestImgLoader):
        if batch_idx >= batch_count:
            break
        start_time = time.time()
        test_loss = test(imgL, imgR, disp_L, save_dir, left_path, disp_path,
                         batch_idx)
        print('Iter (%d / %d) test loss = %.3f, time = %.2f' %
              (batch_idx, len(TestImgLoader), test_loss,
               time.time() - start_time))
        total_test_loss += test_loss

    print('total test loss = %.3f' % (total_test_loss / batch_idx))

if __name__ == "__main__":
    test_model(args.loadmodel,
               './demo_flyingthings3d_flyingthins3dTEST_b2_diff_scale_4', 10)


import sys
import os
# check SBMolGen_PATH setting
if os.getenv('SBMolGen_PATH') == None:
    print("THe SBMolGen_PATH has not defined, please set it before use it!")
    exit(0)
else:
    SBMolGen_PATH=os.getenv('SBMolGen_PATH')
    sys.path.append(SBMolGen_PATH+'/utils')
from subprocess import Popen, PIPE
from math import *
import random
import random as pr
import numpy as np
from copy import deepcopy
import itertools
import time
import math
import argparse
import subprocess
from keras.preprocessing import sequence
from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem import Descriptors
from load_model import loaded_model
from make_smile import zinc_data_with_bracket_original, zinc_processed_with_bracket
from add_node_type_zinc import chem_kn_simulation, make_input_smile,predict_smile,check_node_type,node_to_add,expanded_node
import yaml

class chemical:

    def __init__(self):

        self.position=['&']
        self.num_atom=8
        #self.vl=['\n', '&', 'C', '(', 'c', '1', 'o', '=', 'O', 'N', 'F', '[C@@H]',
        #'n', '-', '#', 'S', 'Cl', '[O-]', '[C@H]', '[NH+]', '[C@]', 's', 'Br', '/', '[nH]', '[NH3+]',
        #'[NH2+]', '[C@@]', '[N+]', '[nH+]', '\\', '[S@]', '[N-]', '[n+]', '[S@@]', '[S-]',
        #'I', '[n-]', 'P', '[OH+]', '[NH-]', '[P@@H]', '[P@@]', '[PH2]', '[P@]', '[P+]', '[S+]',
        #'[o+]', '[CH2-]', '[CH-]', '[SH+]', '[O+]', '[s+]', '[PH+]', '[PH]', '[S@@+]']
        self.vl = ['\n', '&', 'C', '1', 'N', '[C@@H]', '2', '[C@H]', '(', '=', 'O', ')', 'S', 'c', '[S@]', '[nH]', '[O-]', '[N+]', 'n', 'F', '#', '[C@]', '[C@@]', '[S@@]', 'P', '/', '\\', 'Cl', 's', 'Br', 'o', '[NH3+]', 'I', '[n+]', '[nH+]', '3', '[N-]', '[S-]', 'B', '4', '5', '[NH+]', '[Si]', '[P@]', '[NH2+]', '[P@@]', '[N@+]', '6', '[N@@+]', '[S@@+]', '7', '8', '[P@@H]', '[n-]', '[C-]', '[P+]', '[Cu]', '[Ni]', '[Zn]', '[Au-]', '[OH+]']

    def Clone(self):

        st = chemical()
        st.position= self.position[:]
        return st

    def SelectPosition(self,m):
        self.position.append(m)

    def Getatom(self):
        return [i for i in range(self.num_atom)]

class Node:

    def __init__(self, position = None,  parent = None, state = None):
        self.position = position
        self.parentNode = parent
        self.childNodes = []
        self.child=None
        self.wins = 0
        self.visits = 0
        self.nonvisited_atom=state.Getatom()
        self.type_node=[]
        self.depth=0

    def Selectnode(self):

        #s = sorted(self.childNodes, key = lambda c: c.wins/c.visits + 0.8*sqrt(2*log(self.visits)/c.visits))[-1]
        #s=random.choice(self.childNodes)
        ucb=[]
        print('UCB:')
        for i in range(len(self.childNodes)):
            ucb_tmp = self.childNodes[i].wins/self.childNodes[i].visits+ c_val*sqrt(2*log(self.visits)/self.childNodes[i].\
visits)
            ucb.append(ucb_tmp)
            print(self.childNodes[i].position, ucb_tmp,)
        m = np.amax(ucb)
        indices = np.nonzero(ucb == m)[0]
        ind=pr.choice(indices)
        s=self.childNodes[ind]
        print('\n', 'index', ind, self.position, m,)
        return s

    def Addnode(self, m, s):

        n = Node(position = m, parent = self, state = s)
        self.childNodes.append(n)

    def simulation(self,state):
        predicted_smile=predict_smile(model,state)
        input_smile=make_input_smile(predicted_smile)
        logp,valid_smile,all_smile=logp_calculation(input_smile)

        return logp,valid_smile,all_smile

    def Update(self, result):

        self.visits += 1
        self.wins += result

def MCTS(root, verbose = False):

    """initialization of the chemical trees and grammar trees"""
    #run_time=time.time()+3600*48
    start_time = time.time()
    run_time = time.time() + 3600*hours # 3600*24
    rootnode = Node(state = root)
    state = root.Clone()
    """----------------------------------------------------------------------"""

    """global variables used for save valid compounds and simulated compounds"""
    valid_compound=[]
    all_simulated_compound=[]
    desired_compound=[]
    max_logp=[]
    desired_activity=[]
    depth=[]
    min_score=1000
    score_distribution=[]
    min_score_distribution=[]

    generated_dict = {} #dictionary of generated compounds
    dict_id = 1  ## this id used for save best docking pose.
    """----------------------------------------------------------------------"""
    out_f = open(output_dir, 'a')

    while time.time()<=run_time:

        node = rootnode # important !    this node is different with state / node is the tree node
        state = root.Clone() # but this state is the state of the initialization .  too important !!!
        """selection step"""
        node_pool=[]

        while node.childNodes!=[]:
            node = node.Selectnode()
            state.SelectPosition(node.position)
        print("state position:,",state.position)

        if len(state.position)>= 70:
            re= -1.0
            while node != None:
                node.Update(re)
                node = node.parentNode
            continue
        if node.position == '\n':
            re = -1.0
            while node != None:
                node.Update(re)
                node = node.parentNode
            continue

        """------------------------------------------------------------------"""

        """expansion step"""
        expanded=expanded_node(model,state.position,val,loop_num_nodeExpansion)

        new_compound = []
        nodeadded = []
        for n in range(simulation_num):
            nodeadded_tmp = node_to_add(expanded, val)
            all_posible=chem_kn_simulation(model,state.position,val,nodeadded_tmp)
            generate_smile=predict_smile(all_posible,val)
            new_compound_tmp = make_input_smile(generate_smile)
            nodeadded.extend(nodeadded_tmp)
            new_compound.extend(new_compound_tmp)
        print('nodeadded', nodeadded)
        print('new compound', new_compound)
        print('generated_dict', generated_dict)
        print('dict_id', dict_id)
        for comp in new_compound:
            print('lastcomp', comp[-1], ' ... ',comp[-1] == '\n')
        node_index,rdock_score,valid_smile,generated_dict = check_node_type(new_compound, score_type, generated_dict, sa_threshold = sa_threshold, rule = rule5, radical = radical_check, docking_num = docking_num, target_dir = target_dir, hashimoto_filter = hashimoto_filter, dict_id = dict_id, trial = trial)
        valid_compound.extend(valid_smile)
        score_distribution.extend(rdock_score)

        print('node', node_index, 'rdock_score', rdock_score, 'valid', valid_smile)
        #out_f = open(output_dir, 'a')
        #out_f.write(str(valid_smile) + ', '+ str(rdock_score)+', '+str(min_score)+', '+str(len(state.position)))
        out_f.write(str(valid_smile) + ', '+ str(rdock_score)+', '+str(min_score)+', '+str(len(state.position))+', '+str(time.time()-start_time))
        out_f.write('\n')
        out_f.flush()
        #out_f.close()
        dict_id += 1

        if len(node_index)==0:
            re=-1.0
            while node != None:
                node.Update(re)
                node = node.parentNode
        else:
            re_list = []
            #atom_list = [nodeadded[m] for m in node_index]
            atom_checked = []
            for i in range(len(node_index)):
                m=node_index[i]
                atom = nodeadded[m]

                if atom not in atom_checked:
                    node.Addnode(atom, state)
                    node_pool.append(node.childNodes[len(atom_checked)])
                    depth.append(len(state.position))
                    atom_checked.append(atom)
                else:
                    node_pool.append(node.childNodes[atom_checked.index(atom)])

                #node.Addnode(nodeadded[m],state)
                #node.Addnode(nodeadded[m],state)
                #print valid_smile[i], 'node m', m, 'nodeadded[m]', nodeadded[m], 'node.childNodes[i]', node.childNodes[i]
                for child in node.childNodes:
                    print(child.position)
                print('\n')
                #node_pool.append(node.childNodes[i])
                #depth.append(len(state.position))

                score_index = 0 if score_type == 'SCORE' else 1

                print("current minmum score",min_score)
                if rdock_score[i][score_index]<=min_score:
                    min_score_distribution.append(rdock_score[i][score_index])
                    min_score=rdock_score[i][score_index]
                else:
                    min_score_distribution.append(min_score)
                """simulation"""

                if atom == '\n':
                    re = -1
                else:
                    #re=(- (rdock_score[i][score_index] + 20)*0.1)/(1+abs(rdock_score[i][score_index] + 20)*0.1)
                    re=(- (rdock_score[i][score_index] - base_rdock_score)*0.1)/(1+abs(rdock_score[i][score_index] -base_rdock_score)*0.1)
                    #### pj16 reward fuction:
                    #base_rdock_score = -20
                    #reward = (np.tanh(0.1*(abs(rdock_score[max_index])+base_rdock_score)) + 1)/2
                re_list.append(re)
                print('atom', atom, 're_list', re_list)
                #re=(- (rdock_score[i]/100))/(1+abs(rdock_score[i]/100))
                """backpropation step"""

            for i in range(len(node_pool)):

                node=node_pool[i]
                while node != None:
                    node.Update(re_list[i])
                    node = node.parentNode

            for child in node_pool:
                print(child.position, child.wins, child.visits)

    out_f.close()

    """check if found the desired compound"""

    #print "all valid compounds:",valid_compound
    #print "all active compounds:",desired_compound
    print("rdock_score",score_distribution)
    print("num valid_compound:",len(valid_compound))
    print("valid compounds",valid_compound)
    print("depth",depth)
    print("min_score",min_score_distribution)

    return valid_compound

def UCTchemical():
    one_search_start_time=time.time()
    time_out=one_search_start_time+60*10
    state = chemical()
    best = MCTS(root = state,verbose = False)

    return best

if __name__ == "__main__":
    # set parameter
    argvs = sys.argv

    """read yaml file for configuration"""
    f = open(str(argvs[1]), "r+")
    conf = yaml.load(f, Loader=yaml.SafeLoader)
    f.close()

    trial = conf.get('trial', 1)
    c_val = conf.get('c_val', 1.0)
    loop_num_nodeExpansion = conf.get('loop_num_nodeExpansion', 1000)
    target = conf.get('target', 'CDK2')
    target_dir = conf.get('target_path', './')
    hours = conf.get('hours', 1)
    score_type = conf.get('score_type', 'SCORE.INTER') #<SCORE> or <SCORE.INTER>
    docking_num = conf.get('docking_num', 10)
    sa_threshold = conf.get('sa_threshold', 3.5) #if SA > sa_threshold, score = 0. Default sa_threshold = 10
    #RO5: if a compound does not satisfy rule of 5, score = 0.
    rule5 = conf.get('rule5', 1) #0:none, 1: rule of 5, 2: rule of 3
    radical_check = conf.get('radical_check', True)
    simulation_num = conf.get('simulation_num', 3)
    hashimoto_filter = conf.get('hashimoto_filter', True)  # or False, use/not use hashimoto filter
    base_rdock_score = conf.get('base_rdock_score', -20)
    model_name = conf.get('model_name', 'model')

    print('========== display configuration ==========')
    print('trial num is: ', trial)
    print('c_val: ', c_val)
    print('loop_num_nodeExpansion: ', loop_num_nodeExpansion)
    print('target: ', target)
    print('target_dir: ',target_dir)
    print('max run time: ',hours)
    print('score_type: ', score_type)
    print('docking_num: ',docking_num)
    print('sa_threshold: ',sa_threshold)
    print('model_name: ', model_name)
    print('base_rdock_score: ', base_rdock_score)
    print('simulation_num: ',simulation_num)
    print('hashimoto_filter: ', hashimoto_filter)
    """----------------------------------------------------------------------"""

    output_dir = 'result_'+target+'_C'+str(c_val)+'_trial'+str(trial)+'.txt'

    smile_old=zinc_data_with_bracket_original(SBMolGen_PATH + '/data/250k_rndm_zinc_drugs_clean.smi')
    val,smile=zinc_processed_with_bracket(smile_old)
    print('val is ', val)

    out_f = open(output_dir, 'w')
    out_f.write('#valid_smile, rdock_score, min_score, depth, used_time')
    out_f.write('\n')
    out_f.close()

    model=loaded_model(SBMolGen_PATH + '/RNN-model/'+ model_name)  #WM300 not tested
    valid_compound=UCTchemical()


# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmark suite for KPL and feature column implementations."""

import itertools
import math
import random
import string
import time

import numpy as np
import tensorflow.compat.v2 as tf

import keras

class LayerBenchmark(tf.test.Benchmark):
    """Benchmark the layer forward pass."""

    def report(self, name, keras_time, fc_time, iters):
        """Calculate and report benchmark statistics."""
        extras = {
            "fc_avg_time": fc_time,
            "fc_vs_keras_sec": fc_time - keras_time,
            "fc_vs_keras_pct": ((fc_time - keras_time) / fc_time) * 100,
            "keras_faster_ratio": fc_time / keras_time,
        }
        self.report_benchmark(
            iters=iters, wall_time=keras_time, extras=extras, name=name
        )

class StepTimingCallback(keras.callbacks.Callback):
    """A callback that times non-warmup steps of a Keras predict call."""

    def __init__(self):
        self.t0 = None
        self.steps = 0

    def on_predict_batch_begin(self, batch_index, _):
        if batch_index == 2:
            self.t0 = time.time()
        elif batch_index > 2:
            self.steps += 1

    def on_predict_end(self, _):
        self.tn = time.time()
        self.t_avg = (self.tn - self.t0) / self.steps

def create_data(length, num_entries, max_value, dtype):
    """Create a ragged tensor with random data entries."""
    lengths = (np.random.random(size=num_entries) * length).astype(int)
    total_length = np.sum(lengths)
    values = (np.random.random(size=total_length) * max_value).astype(dtype)
    return tf.RaggedTensor.from_row_lengths(values, lengths)

def create_string_data(
    length, num_entries, vocabulary, pct_oov, oov_string="__OOV__"
):
    """Create a ragged tensor with random data entries."""
    lengths = (np.random.random(size=num_entries) * length).astype(int)
    total_length = np.sum(lengths)
    num_oovs = int(pct_oov * total_length)
    values = []
    for _ in range(total_length):
        values.append(random.choice(vocabulary))

    if pct_oov > 0:
        oov_cadence = int(total_length / num_oovs)
        idx = 0
        for _ in range(num_oovs):
            if idx < total_length:
                values[idx] = oov_string
            idx += oov_cadence

    return tf.RaggedTensor.from_row_lengths(values, lengths)

def create_vocabulary(vocab_size):
    base = len(string.ascii_letters)
    n = math.ceil(math.log(vocab_size, base))
    vocab = []
    for i in range(1, n + 1):
        for item in itertools.product(string.ascii_letters, repeat=i):
            if len(vocab) >= vocab_size:
                break
            vocab.append("".join(item))
    return vocab

def run_keras(data, model, batch_size, num_runs, steps_per_repeat=100):
    """Benchmark a Keras model."""
    ds = (
        tf.data.Dataset.from_tensor_slices(data)
        .repeat()
        .prefetch(tf.data.AUTOTUNE)
        .batch(batch_size)
        .cache()
    )
    steps = 0
    times = []
    for _ in range(num_runs):
        steps += steps_per_repeat
        timer = StepTimingCallback()
        # Benchmarked code begins here.
        model.predict(ds, steps=steps, callbacks=[timer])
        # Benchmarked code ends here.
        times.append(timer.t_avg)
    avg_time = np.mean(times)
    return avg_time

def run_fc(data, fc_fn, batch_size, num_runs, steps_per_repeat=100):
    """Benchmark a Feature Column."""

    ds = (
        tf.data.Dataset.from_tensor_slices(data)
        .repeat()
        .prefetch(tf.data.AUTOTUNE)
        .batch(batch_size)
        .cache()
    )

    # Trace the fc_fn
    ds_iter = ds.__iter__()
    fc_fn(next(ds_iter))
    fc_starts = []
    fc_ends = []
    for _ in range(num_runs):
        fc_starts.append(time.time())
        # Benchmarked code begins here.
        for _ in range(steps_per_repeat):
            _ = fc_fn(next(ds_iter))
        # Benchmarked code ends here.
        fc_ends.append(time.time())
    avg_per_step_time = (
        np.array(fc_ends) - np.array(fc_starts)
    ) / steps_per_repeat
    avg_time = np.mean(avg_per_step_time)
    return avg_time


# -*- coding: utf-8 -*-
#
# Copyright (c) 2015 Cisco Systems, Inc. and others.  All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This module contains classifiers for the Active Learning Pipeline.
"""

import os
import logging
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import List, Dict
from sklearn.model_selection import StratifiedKFold
import numpy as np

from .data_loading import LabelMap, DataBucket
from .heuristics import (
    Heuristic,
    KLDivergenceSampling,
    DisagreementSampling,
    EnsembleSampling,
)

from ..components.classifier import Classifier
from ..components.nlp import NaturalLanguageProcessor
from ..constants import (
    TUNE_LEVEL_INTENT,
    ACTIVE_LEARNING_RANDOM_SEED,
    AL_DEFAULT_AGGREGATE_STATISTIC,
    AL_DEFAULT_CLASS_LEVEL_STATISTIC,
    AL_SUPPORTED_AGGREGATE_STATISTICS,
    AL_SUPPORTED_CLASS_LEVEL_STATISTICS,
)
from ..resource_loader import ProcessedQueryList

logger = logging.getLogger(__name__)

MULTI_MODEL_HEURISTICS = (KLDivergenceSampling, DisagreementSampling, EnsembleSampling)

class ALClassifier(ABC):
    """ Abstract class for Active Learning Classifiers."""

    def __init__(self, app_path: str, tuning_level: str):
        """
        Args:
            app_path (str): Path to MindMeld application
            tuning_level (str): The hierarchy level to tune ("domain" or "intent")
        """
        self.app_path = app_path
        self.tuning_level = tuning_level
        self.intent2idx, self.idx2intent, self.domain_indices = self._get_mappings()

    def _get_mappings(self):
        """Get mappings of intents to indices and the indices that map to each domain.

        Returns:
            intent2idx (Dict): Maps intents to indices
            idx2intent (Dict): Maps indices to intents
            domain_indices (Dict): Maps domains to a tuple containing the start and
                ending indexes of intents with the given domain.
        """
        idx = 0
        intent2idx, idx2intent, domain_indices = {}, {}, {}
        for domain in sorted(os.listdir(os.path.join(self.app_path, "domains"))):
            start_idx = idx
            for intent in sorted(
                os.listdir(os.path.join(self.app_path, "domains", domain))
            ):
                intent2idx[f"{domain}.{intent}"] = idx
                idx2intent[idx] = f"{domain}.{intent}"
                idx += 1
            end_idx = idx - 1
            domain_indices[domain] = (start_idx, end_idx)
        return intent2idx, idx2intent, domain_indices

    @abstractmethod
    def train(self):
        raise NotImplementedError(
            "Subclasses must implement their classifier's fit method."
        )

class MindMeldALClassifier(ALClassifier):
    """Active Learning classifier that uses MindMeld classifiers internally.
    Handles the training of MindMeld components (Domain or Intent classifiers)
    and collecting performance statistics (eval_stats)."""

    def __init__(
        self,
        app_path: str,
        tuning_level: str,
        n_classifiers: int,
        aggregate_statistic: str = None,
        class_level_statistic: str = None,
    ):
        """
        Args:
            app_path (str): Path to MindMeld application
            tuning_level (str): The hierarchy level to tune ("domain" or "intent")
            n_classifiers (int): Number of classifiers to be used by multi-model strategies.
        """
        super().__init__(app_path=app_path, tuning_level=tuning_level)
        self.nlp = NaturalLanguageProcessor(self.app_path)
        self.n_classifiers = n_classifiers
        self.aggregate_statistic = MindMeldALClassifier._validate_aggregate_statistic(
            aggregate_statistic
        )
        self.class_level_statistic = (
            MindMeldALClassifier._validate_class_level_statistic(class_level_statistic)
        )

    @staticmethod
    def _validate_aggregate_statistic(aggregate_statistic):
        """Method to validate the aggregate statistic. If an aggregate statistic is not provided
        the default is used. (Options: "accuracy", "f1_weighted", "f1_macro", "f1_micro".)

        Args:
            aggregate_statistic (str): Aggregate statistic to record.
        Returns:
            aggregate_statistic (str): Aggregate statistic to record.
        Raises:
            ValueError: If an invalid value is provided.
        """
        if not aggregate_statistic:
            logger.info(
                "Aggregate statistic not defined, using default: %r.",
                AL_DEFAULT_AGGREGATE_STATISTIC,
            )
            return AL_DEFAULT_AGGREGATE_STATISTIC
        if aggregate_statistic not in AL_SUPPORTED_AGGREGATE_STATISTICS:
            raise ValueError(
                "Not a valid aggregate statistic: {!r}.".format(aggregate_statistic)
            )
        return aggregate_statistic

    @staticmethod
    def _validate_class_level_statistic(class_level_statistic):
        """Method to validate the class-level statistic. If an class-level statistic is not provided
        the default is used. (Options: "f_beta", "percision", "recall")

        Args:
            class_level_statistic (str): Class_level statistic to record.
        Returns:
            class_level_statistic (str): Class_level statistic to record.
        Raises:
            ValueError: If an invalid value is provided.
        """
        if not class_level_statistic:
            logger.info(
                "Class-level statistic not defined, using default: %r.",
                AL_DEFAULT_CLASS_LEVEL_STATISTIC,
            )
            return AL_DEFAULT_CLASS_LEVEL_STATISTIC
        if class_level_statistic not in AL_SUPPORTED_CLASS_LEVEL_STATISTICS:
            raise ValueError(
                "Not a valid class-level statistic: {!r}.".format(class_level_statistic)
            )
        return class_level_statistic

    @staticmethod
    def _get_probs(
        classifier: Classifier, queries: ProcessedQueryList, nlp_component_to_id: Dict
    ):
        """Get the probability distribution for a query across domains or intents.

        Args:
            classifier (MindMeld Classifer): Domain or Intent Classifier
            queries (ProcessedQueryList): List of MindMeld queries
            nlp_component_to_id (Dict): Dictionary mapping domain or intent names to vector index
                positions.

        Returns:
            prob_vector (List[List]]): Probability distribution vectors for given queries.
        """
        queries_prob_vectors = []
        if queries:
            classifier_eval = classifier.evaluate(queries=queries)
            for x in classifier_eval.results:
                query_prob_vector = np.zeros(len(nlp_component_to_id))
                for nlp_component, index in x.probas.items():
                    query_prob_vector[nlp_component_to_id[nlp_component]] = index
                queries_prob_vectors.append(query_prob_vector)
            assert len(queries_prob_vectors) == len(queries)
        return queries_prob_vectors

    def _pad_intent_probs(
        self, ic_queries_prob_vectors: List[List[float]], intents: List
    ):
        """Pads the intent probability array with zeroes for out-of-domain intents.
        Args:
            ic_queries_prob_vectors (List[List[float]]]): 2D Array containing the probability
                distribution for a single query across intents in the query's domain.
            intents (List): List intents in the order that corresponds with the intent
                probabities for the queries. Intents are in the form "domain.intent".
        Returns:
            padded_ic_queries_prob_vectors (List[List[float]]]): 2D Array containing the probability
                distribution for a single query across all intents (including out-of-domain
                intents).
        """
        padded_ic_queries_prob_vectors = []
        for unordered_ic_query_prob_vector in ic_queries_prob_vectors:
            ordered_ic_query_prob_vector = np.zeros(len(self.intent2idx))
            for i, intent in enumerate(intents):
                ordered_ic_query_prob_vector[
                    self.intent2idx[intent]
                ] = unordered_ic_query_prob_vector[i]
            padded_ic_queries_prob_vectors.append(ordered_ic_query_prob_vector)
        return padded_ic_queries_prob_vectors

    def train(self, data_bucket: DataBucket, heuristic: Heuristic):
        """Main training function.

        Args:
            data_bucket (DataBucket): DataBucket for current iteration
            heuristic (Heuristic): Current Heuristic.

        Returns:
            eval_stats (defaultdict): Evaluation metrics to be included in accuracies.json
            confidences_2d (List[List]): 2D array with probability vectors for unsampled queries
            confidences_3d (List[List[List]]]): 3D array with probability vectors for unsampled
                queries from multiple classifiers
            domain_indices (Dict): Maps domains to a tuple containing the start and
                ending indexes of intents with the given domain.
        """
        eval_stats = defaultdict(dict)
        eval_stats["num_sampled"] = len(data_bucket.sampled_queries)
        confidences_2d = self.train_single(data_bucket, eval_stats)
        return_confidences_3d = isinstance(heuristic, MULTI_MODEL_HEURISTICS)
        confidences_3d = (
            self.train_multi(data_bucket) if return_confidences_3d else None
        )
        domain_indices = (
            self.domain_indices if isinstance(heuristic, KLDivergenceSampling) else None
        )
        return eval_stats, confidences_2d, confidences_3d, domain_indices

    def train_single(self, data_bucket: DataBucket, eval_stats: defaultdict = None):
        """Trains a single model to get a 2D probability array for single-model selection strategies.
        Args:
            data_bucket (DataBucket): Databucket for current iteration
            eval_stats (defaultdict): Evaluation metrics to be included in accuracies.json
        Returns:
            confidences_2d (List): 2D array with probability vectors for unsampled queries
        """
        return self._train_single(
            sampled_queries=data_bucket.sampled_queries,
            unsampled_queries=data_bucket.unsampled_queries,
            test_queries=data_bucket.test_queries,
            label_map=data_bucket.label_map,
            eval_stats=eval_stats,
        )

    def _train_single(
        self,
        sampled_queries: ProcessedQueryList,
        unsampled_queries: ProcessedQueryList,
        test_queries: ProcessedQueryList,
        label_map: LabelMap,
        eval_stats: Dict = None,
    ):
        """Helper function to train a single model and obtain a 2D probability array.
        Args:
            sampled_queries (ProcessedQueryList): Current set of sampled queries in DataBucket.
            unsampled_queries (ProcessedQueryList): Current set of unsampled queries in DataBucket.
            test_queries (ProcessedQueryList): Set of test queries in DataBucket.
            label_map LabelMap: Class that stores index mappings for a MindMeld app.
                (Eg. domain2id, domain_to_intent2id)
            eval_stats (Dict): Evaluation metrics to be included in accuracies.json
        Returns:
            confidences_2d (List): 2D array with probability vectors for unsampled queries
        """

        # Domain_Level
        dc_queries_prob_vectors, dc_eval_test = self.domain_classifier_fit_eval(
            sampled_queries=sampled_queries,
            unsampled_queries=unsampled_queries,
            test_queries=test_queries,
            domain2id=label_map.domain2id,
        )
        if eval_stats:
            self._update_eval_stats_domain_level(eval_stats, dc_eval_test)
        confidences_2d = dc_queries_prob_vectors

        # Intent_Level
        if self.tuning_level == TUNE_LEVEL_INTENT:
            (
                ic_queries_prob_vectors,
                ic_eval_test_dict,
            ) = self.intent_classifiers_fit_eval(
                sampled_queries=sampled_queries,
                unsampled_queries=unsampled_queries,
                test_queries=test_queries,
                domain_list=list(label_map.domain2id),
                domain_to_intent2id=label_map.domain_to_intent2id,
            )
            if eval_stats:
                self._update_eval_stats_intent_level(eval_stats, ic_eval_test_dict)
            confidences_2d = ic_queries_prob_vectors

        return confidences_2d

    def train_multi(self, data_bucket: DataBucket):
        """Trains multiple models to get a 3D probability array for multi-model selection strategies.
        Args:
            data_bucket (DataBucket): Databucket for current iteration
        Returns:
            confidences_3d (List[List[List]]]): 3D array with probability vectors for unsampled
                queries from multiple classifiers
        """
        return self._train_multi(
            sampled_queries=data_bucket.sampled_queries,
            unsampled_queries=data_bucket.unsampled_queries,
            test_queries=data_bucket.test_queries,
            label_map=data_bucket.label_map,
        )

    def _train_multi(
        self,
        sampled_queries: ProcessedQueryList,
        unsampled_queries: ProcessedQueryList,
        test_queries: ProcessedQueryList,
        label_map: LabelMap,
    ):
        """Helper function to train multiple models and obtain a 3D probability array.
        Args:
            sampled_queries (ProcessedQueryList): Current set of sampled queries in DataBucket.
            unsampled_queries (ProcessedQueryList): Current set of unsampled queries in DataBucket.
            test_queries (ProcessedQueryList): Set of test queries in DataBucket.
            label_map LabelMap: Class that stores index mappings for a MindMeld app.
                (Eg. domain2Id, domain_to_intent2id)
        Returns:
            confidences_3d (List[List[List]]]): 3D array with probability vectors for unsampled
                queries from multiple classifiers
        """

        sampled_queries_ids = sampled_queries.elements
        skf = StratifiedKFold(
            n_splits=self.n_classifiers,
            shuffle=True,
            random_state=ACTIVE_LEARNING_RANDOM_SEED,
        )
        y = [
            f"{domain}.{intent}"
            for domain, intent in zip(
                sampled_queries.domains(), sampled_queries.intents()
            )
        ]
        fold_sampled_queries_ids = [
            [sampled_queries_ids[i] for i in fold]
            for _, fold in skf.split(sampled_queries_ids, y)
        ]
        fold_sampled_queries_lists = [
            ProcessedQueryList(sampled_queries.cache, fold)
            for fold in fold_sampled_queries_ids
        ]
        confidences_3d = []
        for fold_sample_queries in fold_sampled_queries_lists:
            confidences_3d.append(
                self._train_single(
                    fold_sample_queries, unsampled_queries, test_queries, label_map
                )
            )
        return confidences_3d

    def domain_classifier_fit_eval(
        self,
        sampled_queries: ProcessedQueryList,
        unsampled_queries: ProcessedQueryList,
        test_queries: ProcessedQueryList,
        domain2id: Dict,
    ):
        """Fit and evaluate the domain classifier.
        Args:
            sampled_queries (ProcessedQueryList): List of Sampled Queries
            unsampled_queries (ProcessedQueryList): List of Unsampled Queries
            test_queries (ProcessedQueryList): List of Test Queries
            domain2id (Dict): Dictionary mapping domains to IDs

        Returns:
            dc_queries_prob_vectors (List[List]): List of probability distributions
                for unsampled queries.
            dc_eval_test (mindmeld.models.model.StandardModelEvaluation): Mindmeld evaluation
                object for the domain classifier.
        """
        dc = self.nlp.domain_classifier
        dc.fit(queries=sampled_queries)
        dc_eval_test = dc.evaluate(queries=test_queries)
        dc_queries_prob_vectors = MindMeldALClassifier._get_probs(
            dc, unsampled_queries, domain2id
        )
        return dc_queries_prob_vectors, dc_eval_test

    def _update_eval_stats_domain_level(self, eval_stats: Dict, dc_eval_test):
        """Update the eval_stats dictionary with evaluation metrics from the domain
        classifier.

        Args:
            eval_stats (Dict): Evaluation metrics to be included in accuracies.json
            dc_eval_test (mindmeld.models.model.StandardModelEvaluation): Mindmeld evaluation
                object for the domain classifier.
        """
        eval_stats["accuracies"]["overall"] = dc_eval_test.get_stats()["stats_overall"][
            self.aggregate_statistic
        ]
        logger.info(
            "Overall Domain-level Accuracy: %s", eval_stats["accuracies"]["overall"]
        )

    def intent_classifiers_fit_eval(
        self,
        sampled_queries: ProcessedQueryList,
        unsampled_queries: ProcessedQueryList,
        test_queries: ProcessedQueryList,
        domain_list: Dict,
        domain_to_intent2id: Dict,
    ):
        """Fit and evaluate the domain classifier.
        Args:
            sampled_queries (ProcessedQueryList): List of Sampled Queries.
            unsampled_queries (ProcessedQueryList): List of Unsampled Queries.
            test_queries (ProcessedQueryList): List of Test Queries.
            domain_list (List[str]): List of domains used by the application.
            domain_to_intent2id (Dict): Dictionary mapping intents to IDs.

        Returns:
            ic_queries_prob_vectors (List[List]): List of probability distributions
                for unsampled queries.
            ic_eval_test_dict (Dict): Dictionary mapping a domain (str) to the
                associated ic_eval_test object.
        """
        ic_eval_test_dict = {}
        unsampled_idx_preds_pairs = []
        for domain in domain_list:
            # Filter Queries
            _, filtered_sampled_queries = DataBucket.filter_queries_by_domain(
                sampled_queries, domain
            )
            (
                filtered_unsampled_queries_indices,
                filtered_unsampled_queries,
            ) = DataBucket.filter_queries_by_domain(unsampled_queries, domain)
            _, filtered_test_queries = DataBucket.filter_queries_by_domain(
                test_queries, domain
            )
            # Train
            ic = self.nlp.domains[domain].intent_classifier
            ic.fit(queries=filtered_sampled_queries)
            # Evaluate Test Queries
            ic_eval_test = ic.evaluate(queries=filtered_test_queries)
            if not ic_eval_test:
                unsampled_idx_preds_pairs.extend(
                    (i, np.zeros(len(self.intent2idx)))
                    for i in filtered_unsampled_queries_indices
                )
                continue
            ic_eval_test_dict[domain] = ic_eval_test
            # Get Probability Vectors
            ic_queries_prob_vectors = MindMeldALClassifier._get_probs(
                classifier=ic,
                queries=filtered_unsampled_queries,
                nlp_component_to_id=domain_to_intent2id[domain],
            )
            intents = [
                f"{domain}.{intent}"
                for intent in ic_eval_test.get_stats()["class_labels"]
            ]
            padded_ic_queries_prob_vectors = self._pad_intent_probs(
                ic_queries_prob_vectors, intents
            )
            for i in range(len(filtered_unsampled_queries)):
                unsampled_idx_preds_pairs.append(
                    (
                        filtered_unsampled_queries_indices[i],
                        padded_ic_queries_prob_vectors[i],
                    )
                )
        unsampled_idx_preds_pairs.sort(key=lambda x: x[0])
        padded_ic_queries_prob_vectors = [x[1] for x in unsampled_idx_preds_pairs]
        return padded_ic_queries_prob_vectors, ic_eval_test_dict

    def _update_eval_stats_intent_level(
        self, eval_stats: defaultdict, ic_eval_test_dict: Dict
    ):
        """Update the eval_stats dictionary with evaluation metrics from intent
        classifiers.

        Args:
            eval_stats (defaultdict): Evaluation metrics to be included in accuracies.json.
            ic_eval_test_dict (Dict): Dictionary mapping a domain (str) to the
                associated ic_eval_test object.
        """
        for domain, ic_eval_test in ic_eval_test_dict.items():
            eval_stats["accuracies"][domain] = {
                "overall": ic_eval_test.get_stats()["stats_overall"][
                    self.aggregate_statistic
                ]
            }
            for i, intent in enumerate(ic_eval_test.get_stats()["class_labels"]):
                eval_stats["accuracies"][domain][intent] = {
                    "overall": ic_eval_test.get_stats()["class_stats"][
                        self.class_level_statistic
                    ][i]
                }