
    Di                       U d dl mZ d dlZd dlmZ d dlZd dlZd dlZd dlZd dlm	Z	 d dl
mZ d dlZd dlZd dlZd dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlm Z  d dlm!Z! ejD                  d(d       Z# G d d      Z$ e%e!      Z&dddddgddZ'de(d<   d ejR                  ejT                  jV                  ejT                  jV                   ejT                  jX                  ejT                  jX                    e-d       e-d        e-d      d	Z.d)dZ/d*dZ0d*d Z1d*d!Z2d*d"Z3d*d#Z4d*d$Z5d*d%Z6	 d+	 	 	 	 	 	 	 	 	 	 	 d,d&Z7d-d'Z8y).    )annotationsN)datetime)sleep)Any)JSONSerializable)CategoricalDistribution)FloatDistribution)UpdateFinishedTrialError)BaseStorage)DEFAULT_STUDY_NAME_PREFIX)FrozenStudy)StudyDirection)FrozenTrial)
TrialStatec                     t         N)NotImplementedError     d/home/ubuntu/crypto_trading_bot/.venv/lib/python3.12/site-packages/optuna/testing/pytest_storages.pystorager      s    
r   c                     e Zd Zd<dZd<dZd<dZd<dZd<dZd<dZd<dZ	d<dZ
d<d	Zd<d
Zd<dZd<dZd<dZej"                  j%                  d ej(                          ej(                         f eddd       eddd      fg      	 	 	 	 	 	 	 	 d=d       Zd<dZd<dZd<dZd<dZd<dZd<dZd<dZd<dZd<dZd<dZd<dZ d<dZ!d<d Z"d<d!Z#d<d"Z$d<d#Z%d<d$Z&d<d%Z'ej"                  j%                  d&d'd(gd(d'gg      	 	 	 	 	 	 d>d)       Z(d<d*Z)d<d+Z*d<d,Z+d<d-Z,d<d.Z-ej"                  j%                  d/e.j^                  e.j`                  g      ej"                  j%                  d0g d1d2 e1d3      d4gd2 e1d5      d4g e1d3      d2d4 e1d5      g e1d3      g e1d5      gg      	 	 	 	 	 	 	 	 d?d6              Z2d<d7Z3d<d8Z4d<d9Z5d<d:Z6y;)@StorageTestCasec                   |j                  t        j                  g      }|j                         }t	        |      dk(  sJ |d   j
                  |k(  sJ |d   j                  j                  t              sJ |j                  t        j                  g      }||k7  sJ |j                         }t	        |      dk(  sJ |D ch c]  }|j
                   c}||hk(  sJ t        d |D              sJ y c c}w )N
directions   r      c              3  Z   K   | ]#  }|j                   j                  t               % y wr   )
study_name
startswithr   ).0ss     r   	<genexpr>z8StorageTestCase.test_create_new_study.<locals>.<genexpr>0   s      ^!1<<**+DE^s   )+)
create_new_studyr   MINIMIZEget_all_studieslen	_study_idr    r!   r   all)selfr   study_idfrozen_studies	study_id2r#   s         r   test_create_new_studyz%StorageTestCase.test_create_new_study"   s    ++8O8O7P+Q 002>"a'''a **h666a ++667PQQQ,,9P9P8Q,R	9$$$ 002>"a'''%34498MMMM^~^^^^ 5s   C8c                X   |j                  t        j                  g      }|j                  t        j                  g      }|j                  |       |j                  t        j                  g      }|j	                         }|D ch c]  }|j
                   c}||hk(  sJ y c c}w Nr   r%   r   r&   delete_studyr'   r)   )r+   r   r,   r.   	study_id3r-   r#   s          r   test_create_new_study_unique_idz/StorageTestCase.test_create_new_study_unique_id2   s    ++8O8O7P+Q,,9P9P8Q,R	Y',,9P9P8Q,R	 002%34498MMMM4s   	B'c                l   | j                   j                  }|}|j                  t        j                  g|      }||j                  |      k(  sJ t        j                  t        j                  j                        5  |j                  t        j                  g|       d d d        y # 1 sw Y   y xY w)Nr   r    )test_create_new_study_with_name__name__r%   r   r&   get_study_name_from_idpytestraisesoptuna
exceptionsDuplicatedStudyErrorr+   r   function_namer    r,   s        r   r8   z/StorageTestCase.test_create_new_study_with_name;   s    <<EE"
++&//0Z , 
 W;;HEEEE]]6,,AAB 	b$$1H1H0IV`$a	b 	b 	bs   >#B**B3c                t   |j                  t        j                  g      }|j                  |       |j	                  |      }t        |      dk(  sJ t        j                  t              5  |j                  |dz          d d d        |j                  |       |j                  t        j                  g      }|j	                  |      }t        |      dk(  sJ |j                  |       t        j                  t              5  |j                  |       d d d        y # 1 sw Y   xY w# 1 sw Y   y xY w)Nr   r   r   )
r%   r   r&   create_new_trialget_all_trialsr(   r;   r<   KeyErrorr3   )r+   r   r,   trialss       r   test_delete_studyz!StorageTestCase.test_delete_studyH   s   ++8O8O7P+Q  *''16{a]]8$ 	/  A.	/ 	X&++8O8O7P+Q''16{aX&]]8$ 	+  *	+ 	+	/ 	/	+ 	+s   -D"D."D+.D7c                j   |j                  t        j                  g      }|j                  t        j                  g      }|j                  t        j                  g      }|j                  |       |j	                         D ci c]  }|j
                  | }}||v sJ ||vsJ ||v sJ y c c}w r1   r2   )r+   r   	study_id1r.   r4   r#   studiess          r   /test_delete_study_after_create_multiple_studiesz?StorageTestCase.test_delete_study_after_create_multiple_studies\   s    ,,9P9P8Q,R	,,9P9P8Q,R	,,9P9P8Q,R	Y'+2+B+B+DEa1;;>EEG###'''G### Fs   B0c                   | j                   j                  }|}|j                  t        j                  g|      }|j                  |      |k(  sJ |j                  |      |k(  sJ t        j                  t              5  |j                  d       d d d        t        j                  t              5  |j                  |dz          d d d        y # 1 sw Y   @xY w# 1 sw Y   y xY w)Nr7   z
dummy-namer   )
6test_get_study_id_from_name_and_get_study_name_from_idr9   r%   r   r&   r:   get_study_id_from_namer;   r<   rE   r@   s        r   rM   zFStorageTestCase.test_get_study_id_from_name_and_get_study_name_from_idh   s    SS\\"
++&//0Z , 

 --h7:EEE--j9XEEE ]]8$ 	9**<8	9 ]]8$ 	9**8a<8	9 	9	9 	9	9 	9s    C3CCC&c                   t         j                  ft         j                  ft         j                  t         j                  ft         j                  t         j                  ft         j                  t         j                  gt         j                  t         j                  gfD ][  j                        dfd} |        dz   }t	        j
                  t              5  j                  |       d d d        ] y # 1 sw Y   hxY w)Nr   c                 R    j                        } | t              k(  sJ d       y )NzCDirection of a study should be a tuple of `StudyDirection` objects.)get_study_directionslist)got_directionsr   r,   targets    r   	check_getzDStorageTestCase.test_set_and_get_study_directions.<locals>.check_get   s/    !(!=!=h!G%f5 Y5r   r   )returnNone)r   r&   MAXIMIZEr%   r;   r<   rE   rQ   )r+   r   rU   non_existent_study_idr,   rT   s    `  @@r   !test_set_and_get_study_directionsz1StorageTestCase.test_set_and_get_study_directions{   s    $$&$$&$$n&=&=>$$n&=&=>$$n&=&=>$$n&=&=>
 	DF //6/BH K %-qL!x( D,,-BCD D/	D.D Ds   $DD
	c                   j                  t        j                  g      dfd}t        j	                         D ]  \  }} |||        j                        t        k(  sJ  |dd       dz   }t        j                  t              5  j                  |       d d d        t        j                  t              5  j                  |dd       d d d        y # 1 sw Y   ?xY w# 1 sw Y   y xY w)	Nr   keyvaluec                ^    j                  | |       j                        |    |k(  sJ y r   )set_study_user_attrget_study_user_attrsr\   r]   r   r,   s     r   check_set_and_getzLStorageTestCase.test_set_and_get_study_user_attrs.<locals>.check_set_and_get   s3    ''#u=//9#>%GGGr   datasetImageNetr   r\   strr]   r   rV   rW   )
r%   r   r&   EXAMPLE_ATTRSitemsr`   r;   r<   rE   r_   r+   r   rb   r\   r]   rY   r,   s    `    @r   !test_set_and_get_study_user_attrsz1StorageTestCase.test_set_and_get_study_user_attrs   s    ++8O8O7P+Q	H
 (--/ 	*JCc5)	*++H5FFF 	)Z0 !)1]]8$ 	@(()>?	@ ]]8$ 	O''(=ugN	O 	O		@ 	@	O 	O   C!C-!C*-C6c                    t        |       y r   )-_test_set_and_get_study_user_attrs_for_floatsr+   r   s     r   ,test_set_and_get_study_user_attrs_for_floatsz<StorageTestCase.test_set_and_get_study_user_attrs_for_floats   
    5g>r   c                   j                  t        j                  g      dfd}t        j	                         D ]  \  }} |||        j                        t        k(  sJ  |dd       dz   }t        j                  t              5  j                  |       d d d        t        j                  t              5  j                  |dd       d d d        y # 1 sw Y   ?xY w# 1 sw Y   y xY w)	Nr   r\   r]   c                ^    j                  | |       j                        |    |k(  sJ y r   )set_study_system_attrget_study_system_attrsra   s     r   rb   zNStorageTestCase.test_set_and_get_study_system_attrs.<locals>.check_set_and_get   3    ))(C?11(;C@EIIIr   rc   rd   r   re   )
r%   r   r&   rg   rh   rt   r;   r<   rE   rs   ri   s    `    @r   #test_set_and_get_study_system_attrsz3StorageTestCase.test_set_and_get_study_system_attrs   s    ++8O8O7P+Q	J
 (--/ 	*JCc5)	*--h7=HHH 	)Z0 !)1]]8$ 	B**+@A	B ]]8$ 	Q))*?P	Q 	Q		B 	B	Q 	Qrk   c                    t        |       y r   )/_test_set_and_get_study_system_attrs_for_floatsrn   s     r   .test_set_and_get_study_system_attrs_for_floatsz>StorageTestCase.test_set_and_get_study_system_attrs_for_floats   s
    7@r   c                   |j                  t        j                  g      }t        j	                         D ]  \  }}|j                  |||        |j                  |      t        k(  sJ |j                  |      i k(  sJ |j                  t        j                  g      }t        j	                         D ]  \  }}|j                  |||        |j                  |      t        k(  sJ |j                  |      i k(  sJ y r1   )	r%   r   r&   rg   rh   rs   rt   r`   r_   )r+   r   r,   r\   r]   s        r   *test_study_user_and_system_attrs_confusionz:StorageTestCase.test_study_user_and_system_attrs_confusion   s    ++8O8O7P+Q'--/ 	@JC))(C?	@--h7=HHH++H5;;;++8O8O7P+Q'--/ 	>JC''#u=	>++H5FFF--h72===r   c                   	 	 	 	 	 	 	 	 	 	 	 	 dd}|j                  t        j                  g      }d}t        |      D ]n  }t	        j
                         }t        d       |j                  |      }t        d       t	        j
                         }|j                  |      }	 ||	||||       p t        j                  t              5  |j                  |dz          d d d        |j                  t        j                  g      }
t        |      D ]X  }|j                  |
       |j                  |
      }	|	D ch c]  }|j                   c}t        t        |dz               k(  rXJ  |j                  |      |j                  |
      z   }	t        |	D ch c]  }|j                   c}      d|z  k(  sJ y # 1 sw Y   xY wc c}w c c}w )Nc                   t        |       |dz   k(  sJ t        | D ch c]  }|j                   c}      |dz   k(  sJ | D ch c]  }|j                   c}v sJ | D ch c]  }|j                   c}t        t	        |dz               k(  sJ t        d | D              sJ t        d | D              sJ t        d | D              sJ t        d | D              sJ t        d | D              sJ t        fd| D              sJ t        fd| D              sJ t        d	 | D              sJ t        d
 | D              sJ y c c}w c c}w c c}w )Nr   c              3  V   K   | ]!  }|j                   t        j                  k(   # y wr   )stater   RUNNINGr"   ts     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     Eqww*"4"44E   ')c              3  :   K   | ]  }|j                   i k(    y wr   paramsr   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     6!qxx2~6   c              3  :   K   | ]  }|j                   i k(    y wr   intermediate_valuesr   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     Cqq,,2Cr   c              3  :   K   | ]  }|j                   i k(    y wr   
user_attrsr   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     :aq||r):r   c              3  :   K   | ]  }|j                   i k(    y wr   system_attrsr   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     <q~~+<r   c              3  r   K   | ].  }|j                   k7  r|j                  |j                  k   0 y wr   	_trial_iddatetime_start)r"   r   time_before_creationtrial_ids     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s<      ;;(*q/?/?/K   #77s   47c              3     K   | ]9  }|j                   k(  r(|j                  |j                  cxk  xr k  nc  ; y wr   r   )r"   r   time_after_creationr   r   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   sD      ;;(*q/?/?/K %q'7'7M:MMMs   ?Ac              3  8   K   | ]  }|j                   d u   y wr   datetime_completer   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     Cqq**d2C   c              3  8   K   | ]  }|j                   d u   y wr   r]   r   s     r   r$   zOStorageTestCase.test_create_new_trial.<locals>._check_trials.<locals>.<genexpr>   s     71qww$7r   r(   r   numbersetranger*   )rF   idxr   r   r   r   s     ``` r   _check_trialsz<StorageTestCase.test_create_new_trial.<locals>._check_trials   sa    v;#')))V445q@@@V<<<<<&,-AHH-U37^1DDDDEfEEEE6v6666CFCCCC:6::::<V<<<<    
     
 CFCCCC77777' 5<-s   EE$Er      MbP?r   r   )rF   list[FrozenTrial]r   intr   r   r   r   r   r   rV   rW   )r%   r   r&   r   r   nowr   rC   rD   r;   r<   rE   r   r   r(   r   )r+   r   r   r,   n_trial_in_studyir   r   r   rF   r.   r   s               r   test_create_new_trialz%StorageTestCase.test_create_new_trial   s   	8%	8	8 	8 #+		8
 "*	8 	8: ++8O8O7P+Q'( 	ZA#+<<> %L//9H%L"*,,.++H5F&!X/CEXY	Z ]]8$ 	3$$X\2	3 ,,9P9P8Q,R	'( 	CA$$Y/++I6F&,-AHH-U1q5\1BBBB	C ''1G4J4J94UU0AAKK01Q9I5IIII	3 	3 . 1s   	F<;GG<Gzstart_time,complete_timei  	   r   r   c                   t        t        j                  d||ddidt        dd      iddidd	id
ddddd      dfd}|j	                  t
        j                  g      }d}t        |      D ]0  }|j                  |      }|j                  |      }	 ||	||       2 t        j                  t              5  |j                  |dz          d d d        |j	                  t
        j                  g      }
t        |      D ]Z  }|j                  |
       |j                  |
      }	|	D ch c]  }|j                   c}t        t        |dz               k(  rZJ  |j                  |      |j                  |
      z   }	t        |	D ch c]  }|j                    c}      d|z  k(  sJ y # 1 sw Y   xY wc c}w c c}w )Ni'  x      ?r   r   foobarbaz{   
   d   i  )r   r   r   7   )r   r]   r   r   r   distributionsr   r   r   r   r   c                   t        |       |dz   k(  sJ t        | D ch c]  }|j                   c}      |dz   k(  sJ || D ch c]  }|j                   c}v sJ | D ch c]  }|j                   c}t        t	        |dz               k(  sJ t        fd| D              sJ t        fd| D              sJ t        fd| D              sJ t        fd| D              sJ t        fd| D              sJ t        fd| D              sJ t        fd| D              sJ t        fd	| D              sJ t        fd
| D              sJ y c c}w c c}w c c}w )Nr   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r"   r   template_trials     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>-        G1qww."6"66G   #&c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>.  s      IQqxx>#8#88Ir   c              3  P   K   | ]  }|j                   j                   k(    y wr   )r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>/  s      W1q.*F*FFWr   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>0  s#     cWXq,,0R0RRcr   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>1  s      QQq||~'@'@@Qr   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>2  s      Uq~~)D)DDUr   c              3  P   K   | ]  }|j                   j                   k(    y wr   )r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>3  s"     YQq''>+H+HHYr   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>4  s#     _STq**n.N.NN_r   c              3  P   K   | ]  }|j                   j                   k(    y wr   r   r   s     r   r$   zcStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials.<locals>.<genexpr>5  r   r   r   )rF   r   r   r   r   s       r   r   zPStorageTestCase.test_create_new_trial_with_template_trial.<locals>._check_trials(  sI   v;#')))V445q@@@V<<<<<&,-AHH-U37^1DDDDGGGGGI&IIIIWPVWWWWc\bccccQ&QQQQUfUUUUYRXYYYY_X^____GGGGG 5<-s   EE"E#r   r   r   r   )rF   r   r   r   r   r   rV   rW   )r   r   COMPLETEr	   r%   r   r&   r   rC   rD   r;   r<   rE   r   r   r(   r   )r+   r   
start_timecomplete_timer   r,   r   r   r   rF   r.   r   r   s               @r   )test_create_new_trial_with_template_trialz9StorageTestCase.test_create_new_trial_with_template_trial  s    %%%%+: 1!Q 78u~$&34 8
	H ++8O8O7P+Q'( 	/A///XH++H5F&!X.	/ ]]8$ 	3$$X\2	3 ,,9P9P8Q,R	'( 	CA$$Y~$N++I6F&,-AHH-U1q5\1BBBB	C
 ''1G4J4J94UU0AAKK01Q9I5IIII	3 	3 . 1s   <F10F=G1F:c                f   |j                  t        j                  g      }|j                  |      }|j	                  |      dk(  sJ |j                  |      }|j	                  |      dk(  sJ t        j                  t              5  |j	                  |dz          d d d        y # 1 sw Y   y xY w)Nr   r   r   )r%   r   r&   rC   get_trial_number_from_idr;   r<   rE   r+   r   r,   r   s       r   test_get_trial_number_from_idz-StorageTestCase.test_get_trial_number_from_idM  s    ++8O8O7P+Q++H5//9Q>>>++H5//9Q>>>]]8$ 	;,,X\:	; 	; 	;s   	B''B0c                   |j                  t        j                  g      }t        D cg c]  }|j	                  |       }}t        |t              D ](  \  }}|t        j                  k(  r|j                  |      j                  t        j                  k(  sJ |j                  |      j                  }|j                  |||j                         rdnd        |j                  |      j                  |k(  sJ |t        j                  k(  r |j                  |      j                  |k(  sJ |j                         r|j                  |      j                  
J |j                  |      j                  )J  t        j                   t"              5  t%        |      dz   }|j                  |t        j&                         d d d        t        D ]  }|j                         s|j	                  |      }|j                  ||d       t        D ]P  }	d }
|	t        j&                  k(  rdg}
t        j                   t(              5  |j                  ||	|
       d d d        R  y c c}w # 1 sw Y   xY w# 1 sw Y   pxY w)Nr   )        r   valuesr   r   r   )r%   r   r&   
ALL_STATESrC   zipr   WAITING	get_trialr   r   r   set_trial_state_valuesis_finishedr   r;   r<   rE   maxr   r
   )r+   r   r,   _	trial_idsr   r   datetime_start_prevnon_existent_trial_idstate2r   s              r   %test_set_trial_state_values_for_statez5StorageTestCase.test_set_trial_state_values_for_stateZ  sm   ++8O8O7P+QAKLAW--h7L	L"9j9 	MOHe
***$$X.44
8J8JJJJ")"3"3H"="L"L**8I8I8KfQU +  $$X.44===
***((2AAEXXXX  "((2DDPPP((2DDLLL	M$ ]]8$ 	$'	NQ$6!**% )) + 	   	ZE$$&//9H**85*P$ Z-1Z000!UF]]#;< Z2286RX2YZ ZZ	Z7 M(	 	"Z Zs   I,0I1
I=1I:=Jc                    t        |       y r   )'_test_set_trial_state_values_for_floatsrn   s     r   &test_set_trial_state_values_for_floatsz6StorageTestCase.test_set_trial_state_values_for_floats  s
    /8r   c                T   t        |ddd      \  }}|j                         D ]  \  }}|j                         D ]  \  }}|j                  |      |j                  k(  sJ |j                  j	                         D ]B  }|j                  ||      |j                  |   j                  |j                  |         k(  rBJ    t        d |j                         D              dz   }t        j                  t              5  |j                  |       d d d        t        j                  t              5  |j                  |d       d d d        |dz
  }	t        j                  t              5  |j                  |	d       d d d        y # 1 sw Y   wxY w# 1 sw Y   OxY w# 1 sw Y   y xY w)Nr      r   n_studyn_trialseedc              3  J   K   | ]  }|j                         D ]  }|   y wr   keysr"   tstids      r   r$   zLStorageTestCase.test_get_trial_param_and_get_trial_params.<locals>.<genexpr>  "     L"'')L3LL   !#paramAz	dummy-key)_setup_studiesrh   get_trial_paramsr   r   get_trial_paramr   to_internal_reprr   r   r;   r<   rE   )
r+   r   r   study_to_trialstrial_id_to_trialr   expected_trialr\   r   existent_trial_ids
             r   )test_get_trial_param_and_get_trial_paramsz9StorageTestCase.test_get_trial_param_and_get_trial_params  s   +GQPQR?$3$9$9$; 	C A ,=,C,C,E C(.//9^=R=RRRR)00557 CC"228SA^EaEaF&&~'<'<S'ABC C CCC	C Lo446LLqP 	 ]]8$ 	<$$%:;	<]]8$ 	E##$98D	E1A5]]8$ 	D##$5{C	D 	D	< 	<	E 	E	D 	Ds$   >F1F*FFFF'c           	        |j                  t        j                  g      }|j                  |      }|j                  |      }|j                  |j                  t        j                  g            }t	        dd      }t        d      }t        d      }t	        ddd	
      }	|j                  |dd|       |j                  |dd|       |j                  |d      dk(  sJ |j                  |d      dk(  sJ |j                  |      j                  dddk(  sJ |j                  |      dddk(  sJ |j                  |dd|       |j                  |dd|	       |j                  |d      dk(  sJ |j                  |d      dk(  sJ |j                  |      j                  dddk(  sJ |j                  |      dddk(  sJ t        j                  t              5  |j                  |dd|	       d d d        t        j                  t              5  |j                  |ddt        d             d d d        |j                  |t        j                   dg       t        j                  t"              5  |j                  |dd|       d d d        t        j                  t$              5  |j                  |d       d d d        t        j                  t"              5  |j                  |dd|	       d d d        |j                  |dd|       |j                  |d      dk(  sJ |j                  |      j                  ddik(  sJ |j                  |      ddik(  sJ t'        |||g      dz   }
t        j                  t$              5  |j                  |
dd|       d d d        y # 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   ;xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   y xY w)Nr         ?       @lowhigh)ShibuyaEbisuMegurochoices)r  Shinseng      Y@T)r  r  logr   r   yr   r  )r   r  333333?z皙?)r   r  )r  r  r  r   r   皙?r   r	  )r%   r   r&   rC   r	   r   set_trial_paramr   r   r   r   r;   r<   
ValueErrorr   r   r   r
   rE   r   )r+   r   r,   
trial_id_1
trial_id_2
trial_id_3distribution_xdistribution_y_1distribution_y_2distribution_zr   s              r   test_set_trial_paramz$StorageTestCase.test_set_trial_param  s   ++8O8O7P+Q--h7
--h7
--$$1H1H0I$J


 +s=2;YZ2;QR*sDI 	
CnE
C4DE&&z373>>>&&z371<<<  ,33Sx7PPPP''
3Sx7PPPP 	
CnE
CnE&&z373>>>&&z373>>>  ,33Ss7KKKK''
3Ss7KKKK ]]:& 	J##JS.I	J ]]:& 	##C$;Db$c	
 	&&z9L9LVYUZ&[]]34 	J##JQ8HI	J ]]8$ 	5##J4	5 ]]34 	J##JS.I	J 	
C4DE&&z371<<<  ,33Y7GGGG''
3Y7GGGG !$ZZ$H IA M]]8$ 	U##$93^T	U 	U7	J 	J	 		J 	J	5 	5	J 	J	U 	UsH   $N1N>=O3O'O%O11N;>OOO"%O.1O:c                    t        |       y r   )(_test_set_and_get_trial_param_for_floatsrn   s     r   'test_set_and_get_trial_param_for_floatsz7StorageTestCase.test_set_and_get_trial_param_for_floats  s
    09r   c                   |j                  t        j                  g      }|j                  |      }|j                  |      }|j                  |j                  t        j                  g            }|j                  |      }|j                  |      }|j	                  |t
        j                  d       |j	                  |t
        j                  t        d      f       |j	                  |t
        j                  d       |j	                  |t
        j                  g d       |j                  |      j                  dk(  sJ |j                  |      j                  J |j                  |      j                  t        d      k(  sJ |j                  |      j                  g dk(  sJ |j                  |      j                  g dk(  sJ t        |||||      dz   }t        j                  t              5  |j	                  |t
        j                  d       d d d        t        j                  t               5  |j	                  |t
        j                  d       d d d        y # 1 sw Y   NxY w# 1 sw Y   y xY w)	Nr   )r   r   inf)r  皙?r  r   r   )r   )r%   r   r&   rC   r   r   r   floatr   r   r]   r   r   r;   r<   rE   r
   )	r+   r   r,   r  r  r  
trial_id_4
trial_id_5r   s	            r   &test_set_trial_state_values_for_valuesz6StorageTestCase.test_set_trial_state_values_for_values  sV   ++8O8O7P+Q--h7
--h7
--$$1H1H0I$J

 --h7
--h7
 	&&z9L9LU[&\&&j115</ 	' 	
 	&&j00 	' 	
 	&&j00 	' 	
   ,22c999  ,22:::  ,22eElBBB  ,33FFF  ,33FFF #J
J
T^ _bc c]]8$ 	**%Z-@-@ + 	 ]]34 	_**:Z=P=PY]*^	_ 	_	 		_ 	_s   ?#I/#I;/I8;Jc                Z   |j                  t        j                  g      }|j                  |      }|j                  |      }|j                  |j                  t        j                  g            }|j                  |      }|j	                  |dd       |j	                  |dd       |j	                  |dd       |j	                  |dd       |j	                  |dd       |j	                  |d	t        d
             |j	                  |dt        d             |j                  |      j                  dddk(  sJ |j                  |      j                  i k(  sJ |j                  |      j                  dddt        d
      dk(  sJ t        j                  |j                  |      j                  d         sJ |j	                  |dd       |j                  |      j                  dddk(  sJ t        ||||      dz   }t        j                  t              5  |j	                  |dd       d d d        |j                  |t        j                   dg       t        j                  t"              5  |j	                  |dd       d d d        y # 1 sw Y   bxY w# 1 sw Y   y xY w)Nr   r   r  r   r  r  r   r   r   r  nan)r   r   )r   r   r   r   r  r   r   )r%   r   r&   rC   set_trial_intermediate_valuer   r   r   npisnanr   r;   r<   rE   r   r   r   r
   )r+   r   r,   r  r  r  r!  r   s           r   !test_set_trial_intermediate_valuez1StorageTestCase.test_set_trial_intermediate_value  s   ++8O8O7P+Q--h7
--h7
--$$1H1H0I$J

 --h7
 	,,ZC@,,ZC@,,ZC@,,ZC@,,ZC@,,ZE%LI,,ZE%LI  ,@@PSDTTTT  ,@@BFFF  ,@@U|	E
 
 	
 
 xx))*5II!LMMM 	,,ZC@  ,@@PSDTTTT #J
J
 SVW W]]8$ 	P001F3O	P 	&&z9L9LVYUZ&[]]34 	E00QD	E 	E	P 	P
	E 	Es    J8J!J!J*c                    t        |       y r   )-_test_set_trial_intermediate_value_for_floatsrn   s     r   ,test_set_trial_intermediate_value_for_floatsz<StorageTestCase.test_set_trial_intermediate_value_for_floats0  rp   r   c                6   t        ddd      \  }}t        fd|j                         D              sJ t        d |j                         D              dz   }t	        j
                  t              5  j                  |       d d d        y # 1 sw Y   y xY w)Nr   r   r   r   c              3     K   | ]:  }|j                         D ]%  \  }}j                  |      |j                  k(   ' < y wr   )rh   get_trial_user_attrsr   r"   rF   r   trialr   s       r   r$   z<StorageTestCase.test_get_trial_user_attrs.<locals>.<genexpr>5  sN      
#)<<>
  % ((2e6F6FF
F
   A Ac              3  J   K   | ]  }|j                         D ]  }|   y wr   r   r   s      r   r$   z<StorageTestCase.test_get_trial_user_attrs.<locals>.<genexpr>;  &      ]SUSZSZS\ ]C ] ]r   r   )r   r*   r   r   r;   r<   rE   r/  r+   r   r   r   non_existent_trials    `   r   test_get_trial_user_attrsz)StorageTestCase.test_get_trial_user_attrs3  s    +GQPRS? 
)002
 
 	
 
 ! ]/E/E/G ]]`aa]]8$ 	=(();<	= 	= 	=   4BBc                   j                  j                  t        j                  g            }dfd}t        j                         D ]  \  }} ||||        j                  |      j                  t        k(  sJ  ||dd       j                  j                  t        j                  g            } ||dd       t        j                  |      j                        d	k(  sJ j                  |      j                  d   dk(  sJ t        ||h      d	z   }t        j                  t              5  j                  |dd       d d d        j                  |t        j                   d
g       t        j                  t"              5  j                  |dd       d d d        y # 1 sw Y   bxY w# 1 sw Y   y xY w)Nr   r\   r]   c                r    j                  | ||       j                  |       j                  |   |k(  sJ y r   )set_trial_user_attrr   r   r   r\   r]   r   s      r   rb   zCStorageTestCase.test_set_trial_user_attr.<locals>.check_set_and_getD  s9    ''#u=$$X.99#>%GGGr   rc   rd   baseline_scorer   r   r   r   r   r   r\   rf   r]   r   rV   rW   )rC   r%   r   r&   rg   rh   r   r   r(   r   r;   r<   rE   r;  r   r   r   r
   )r+   r   r  rb   r\   r]   r  r   s    `      r   test_set_trial_user_attrz(StorageTestCase.test_set_trial_user_attr?  s   --$$1H1H0I$J

	H
 (--/ 	6JCj#u5	6  ,77=HHH 	*i< --$$1H1H0I$J

 	*&6>7$$Z0;;<AAA  ,778HIURRR !$Z$< = A]]8$ 	O''(=ugN	O 	&&z9L9LVYUZ&[]]34 	D''
E7C	D 	D	O 	O
	D 	Ds   ;F0F<0F9<Gc                    t        |       y r   ),_test_set_and_get_trial_user_attr_for_floatsrn   s     r   +test_set_and_get_trial_user_attr_for_floatsz;StorageTestCase.test_set_and_get_trial_user_attr_for_floatsb  s
    4W=r   c                6   t        ddd      \  }}t        fd|j                         D              sJ t        d |j                         D              dz   }t	        j
                  t              5  j                  |       d d d        y # 1 sw Y   y xY w)Nr   r   r   r   c              3     K   | ]:  }|j                         D ]%  \  }}j                  |      |j                  k(   ' < y wr   )rh   get_trial_system_attrsr   r0  s       r   r$   z>StorageTestCase.test_get_trial_system_attrs.<locals>.<genexpr>g  sN      
#)<<>
  % **848J8JJ
J
r2  c              3  J   K   | ]  }|j                         D ]  }|   y wr   r   r   s      r   r$   z>StorageTestCase.test_get_trial_system_attrs.<locals>.<genexpr>m  r4  r   r   )r   r*   r   r   r;   r<   rE   rE  r5  s    `   r   test_get_trial_system_attrsz+StorageTestCase.test_get_trial_system_attrse  s    +GQPRS? 
)002
 
 	
 
 ! ]/E/E/G ]]`aa]]8$ 	?**+=>	? 	? 	?r8  c                
   j                  t        j                  g      }j                  |      }dfd}t        j                         D ]  \  }} ||||        j                  |      j                  }|t        k(  sJ  ||dd       j                  |      } ||dd       j                  |      j                  }|ddik(  sJ t        ||h      d	z   }	t        j                  t              5  j                  |	dd       d d d        j                  |t        j                  d
g       t        j                  t               5  j                  |dd       d d d        y # 1 sw Y   bxY w# 1 sw Y   y xY w)Nr   r\   r]   c                ^    j                  | ||       j                  |       |   |k(  sJ y r   )set_trial_system_attrrE  r<  s      r   rb   zEStorageTestCase.test_set_trial_system_attr.<locals>.check_set_and_getu  ru   r   rc   rd   r=  r   r   r   r   r>  )r%   r   r&   rC   rg   rh   r   r   r   r;   r<   rE   rJ  r   r   r   r
   )
r+   r   r,   r  rb   r\   r]   r   r  r   s
    `        r   test_set_trial_system_attrz*StorageTestCase.test_set_trial_system_attrq  s   ++8O8O7P+Q--h7
	J
 (--/ 	6JCj#u5	6((4AA},,, 	*i< --h7
*&6>((4AA 0%8888 !$Z$< = A]]8$ 	Q))*?P	Q 	&&z9L9LVYUZ&[]]34 	F))*eWE	F 	F	Q 	Q
	F 	Fs   8E-E9-E69Fc                    t        |       y r   )._test_set_and_get_trial_system_attr_for_floatsrn   s     r   -test_set_and_get_trial_system_attr_for_floatsz=StorageTestCase.test_set_and_get_trial_system_attr_for_floats  s
    6w?r   c                   t        |ddd      \  }}|j                         }t        |      t        |      k(  sJ |j                         D ]  \  }}d }|D ]  }|j                  |j                  k(  s|} n |J |j
                  |j
                  k(  sJ |j                  |j                  k(  sJ |j                  |j                  k(  sJ |j                  |j                  k(  rJ  y )Nr   .   r   )r   r'   r(   rh   r    	directionr   r   )r+   r   expected_frozen_studiesr   r-   expected_frozen_studyfrozen_studyr#   s           r   test_get_all_studiesz$StorageTestCase.test_get_all_studies  s
   %3GRQSZ\%]" 002>"c*A&BBBB(?(E(E(G 
	S$A$/3L# <<#8#C#CC#$L  +++))-B-L-LLLL**.C.N.NNNN**.C.N.NNNN,,0E0R0RRRR
	Sr   c                   t        |ddd      \  }}|j                         D ]<  \  }}|j                         D ]$  }|j                  |j                        }||k(  r$J  > t        d |j                         D              dz   }t        j                  t              5  |j                  |       d d d        y # 1 sw Y   y xY w)Nr      /   r   c              3  J   K   | ]  }|j                         D ]  }|   y wr   r   r   s      r   r$   z1StorageTestCase.test_get_trial.<locals>.<genexpr>  r   r   r   )	r   rh   r   r   r   r   r;   r<   rE   )r+   r   r   r   expected_trialsr   r1  r   s           r   test_get_trialzStorageTestCase.test_get_trial  s    +GQQST?"1"7"7"9 	/A"1"8"8": /)).*B*BC.../	/ Lo446LLqP 	 ]]8$ 	534	5 	5 	5s   B99Cc                j   t        |ddd      \  }}|j                         D ]3  \  }}|j                  |      }|D ]  }||j                     }||k(  rJ  5 t	        |j                               dz   }	t        j                  t              5  |j                  |	       d d d        y # 1 sw Y   y xY w)Nr   rW  0   r   r   )	r   rh   rD   r   r   r   r;   r<   rE   )
r+   r   r   r   r,   rZ  rF   r1  r   rY   s
             r   test_get_all_trialsz#StorageTestCase.test_get_all_trials  s    +GQQST?)8)>)>)@ 	/%Ho++H5F /!0!A.../	/ !$O$8$8$: ;a ?]]8$ 	:""#89	: 	: 	:s   B))B2param_namesabc           
        |j                  t        j                  g      }|j                  |t        j
                  j                  t        j                              }|D ]!  }|j                  ||dt        dd             # |j                  |      }t        |d   j                  j                               |k(  sJ y )Nr   r   r   r   r   )distributionr   )r%   r   r&   rC   r=   r1  create_trialr   r   r  r	   rD   rR   r   r   )r+   r   r_  r,   r   
param_namerF   s          r    test_get_all_trials_params_orderz0StorageTestCase.test_get_all_trials_params_order  s     ++8O8O7P+Q++fll//j6H6H/I
 & 	J##*c8I#s8S $ 	
 ''1F1I$$))+,;;;r   c                   t        |ddd      \  }}|D ]r  }|j                  |d      }t        |      t        ||         k(  sJ t        j                  |      }d|d   j
                  d	<   |j                  |d
      }||k(  rrJ  y )Nr   r   1   r   Tdeepcopyr  r   r   F)r   rD   r(   copyrj  r   )r+   r   r-   r   r,   trials0trials0_originaltrials1s           r   #test_get_all_trials_deepcopy_optionz3StorageTestCase.test_get_all_trials_deepcopy_option  s    *8!UV]_*`'& 		/H,,X,EGw<3x'@#AAAA  $}}W5%(GAJc",,X,FG#w...		/r   c                   |j                  t        j                  g      }t        j                  d      }t
        j                  t
        j                  t
        j                  f}|D ]N  }t        |      }||_	        |j                  t
        j                  k(  rdgnd |_
        |j                  ||       P |j                  |d       }t        |      dk(  sJ |j                  |t
        j                  f      }t        |      dk(  sJ t        d |D              sJ |j                  |t
        j                  t
        j                  f      }t        |      dk(  sJ t        d	 |D              sJ |j                  |d
      }t        |      dk(  sJ t        D cg c],  }|t
        j                  k7  s|t
        j                  k7  s+|. }	}|	D ]&  }|j                  ||f      }t        |      dk(  r&J  y c c}w )Nr   3   r   r   )statesr   r   c              3  V   K   | ]!  }|j                   t        j                  k(   # y wr   )r   r   r   r   s     r   r$   zCStorageTestCase.test_get_all_trials_state_option.<locals>.<genexpr>  s     Ba177j111Br   c              3  t   K   | ]0  }|j                   t        j                  t        j                  fv  2 y wr   )r   r   r   PRUNEDr   s     r   r$   zCStorageTestCase.test_get_all_trials_state_option.<locals>.<genexpr>  s*     W1177z22J4E4EFFWs   68r   r   )r%   r   rX   randomRandomr   r   ru  _generate_trialr   r   rC   rD   r(   r*   r   )
r+   r   r,   	generatorrr  r   r   rF   r#   other_statess
             r    test_get_all_trials_state_optionz0StorageTestCase.test_get_all_trials_state_option  s   ++8O8O7P+QMM"%	 
  	AE	*AAG !:+>+> >uDAH$$Xa$@		A '''>6{a''*:M:M9O'P6{aB6BBBB''*:M:MzO`O`9a'b6{aWPVWWWW'''<6{a "
Q**=*=%=!zGXGXBXA
 
 " 	$E++HeX+FFv;!###	$
s   "H;HHc                   t        |ddd      \  }}|j                         D ]  }|j                  |d      }t        j                  |      }|D ]  }|j
                  j                         sb|j                  |j                  ddt        dd	             |j                  |j                  d
d       |j                  |j                  dd       |j
                  t        j                  k(  r|j                  dz  dk(  r-|j                  |j                  t        j                   dg       |j                  dz  d	k(  rK|j#                  |j                  dd       |j                  |j                  t        j$                  dg       9|j                  |j                  t        j&                         e|j
                  t        j(                  k(  s|j                  |j                  t        j                          ||k(  rJ  y )Nr   rW  r]  r   Fri  paramXr   r   	usr_attrX	sys_attrXr   )r   r   rD   rk  rj  r   r   r  r   r	   r;  rJ  r   r   r   r   r   r&  ru  FAILr   )r+   r   r   r   r,   rF   deepcopied_trialsr1  s           r    test_get_all_trials_not_modifiedz0StorageTestCase.test_get_all_trials_not_modified  s   +GQQST?',,. 	/H++Hu+EF $f 5 X{{..0++EOOXqJ[\]_`Jab//aP11%//;PQR;;*"4"44||a'1,66u
H[H[^_]`a)Q.<<U__aQRS66u
HYHY\][^_66u
X[[J$6$66225??JDVDVWX" ....+	/r   c                   t        |ddd      \  }}|D ]  }|j                  |      dk(  rJ  t        |j                               dz   }t	        j
                  t              5  |j                  |      sJ 	 d d d        y # 1 sw Y   y xY w)Nr      2   r   r   )r   get_n_trialsr   r   r;   r<   rE   )r+   r   study_id_to_frozen_studiesr   r,   rY   s         r   test_get_n_trialsz!StorageTestCase.test_get_n_trials  s    (6wST[](^%"A2 	7H''1Q666	7 !$$>$C$C$E F J]]8$ 	?''(=>>>	? 	? 	?s   %BBc                   |j                  t        j                  f      }t        j                  d      }t
        j                  t
        j                  t
        j                  g}|D ]N  }t        |      }||_	        |j                  t
        j                  k(  rdgnd |_
        |j                  ||       P |j                  |t
        j                        dk(  sJ |j                  |t
        j                        dk(  sJ t        D cg c],  }|t
        j                  k7  s|t
        j                  k7  s+|. }}|D ]  }|j                  ||      dk(  rJ  y c c}w )Nr   rq  r   r   r   r   r   )r%   r   rX   rv  rw  r   r   ru  rx  r   r   rC   r  r   )r+   r   r,   ry  rr  r#   r   rz  s           r   test_get_n_trials_state_optionz.StorageTestCase.test_get_n_trials_state_option&  sN   ++8O8O7Q+RMM"%	 
  	AA	*AAG !:+>+> >uDAH$$Xa$@		A ##Hj.A.ABaGGG##Hj.?.?@AEEE "
Q**=*=%=!zGXGXBXA
 
  	:A''!4999	:
s   E$E8ErQ  r   )r   r   r   r   r  r   z-infc                j   |j                  |g      }t        j                  t              5  |j	                  |       d d d        t        j                  t
              5  |j	                  |dz          d d d        t        j                  d      }|D ]<  }t        |      }t        j                  |_        ||_        |j                  ||       > |t        j                  k(  rt!        |      n
t#        |      }|j	                  |      j                  |k(  sJ y # 1 sw Y   xY w# 1 sw Y   xY w)Nr   r   rq  r   )r%   r;   r<   r  get_best_trialrE   rv  rw  rx  r   r   r   r]   rC   r   rX   r   min)	r+   r   rQ  r   r,   ry  vr   expected_values	            r   test_get_best_trialz#StorageTestCase.test_get_best_trial?  s    ++	{+C]]:& 	-""8,	- ]]8$ 	1""8a<0	1 MM"%	 	NA,Y7N#-#6#6N #$N $$Xn$M		N
 )2^5L5L(LVRUV\R]%%h/55GGG	- 	-	1 	1s   D D)D&)D2c                   |j                  t        j                  t        j                  f      }t	        j
                  d      }t        d      D ]A  }t        |      }t        j                  |_
        ||dz   g|_        |j                  ||       C t        j                  t              5  |j!                  |       d d d        y # 1 sw Y   y xY w)Nr   rq  r   r   r   )r%   r   rX   r&   rv  rw  r   rx  r   r   r   r   rC   r;   r<   RuntimeErrorr  )r+   r   r,   ry  r   r   s         r   4test_get_best_trial_for_multi_objective_optimizationzDStorageTestCase.test_get_best_trial_for_multi_objective_optimization_  s    ++&//1H1HI , 
 MM"%	q 	NA,Y7N#-#6#6N %&AJN!$$Xn$M		N ]]<( 	-""8,	- 	- 	-s   .C		Cc                0   t        j                  t              5  |j                  dd       d d d        |j	                  t
        j                  g      }t        j                  t              5  |j                  |d       d d d        |j                  |      }||j                  |d      k(  sJ |j	                  t
        j                  g      }|j                  |      }||j                  |d      k(  sJ y # 1 sw Y   xY w# 1 sw Y   xY w)Nr   )r,   trial_numberr   )r  )r;   r<   rE   'get_trial_id_from_study_id_trial_numberr%   r   r&   rC   r   s       r   ,test_get_trial_id_from_study_id_trial_numberz<StorageTestCase.test_get_trial_id_from_study_id_trial_numbern  s&   ]]8$ 	X;;QUV;W	X ++8O8O7P+Q]]8$ 	V;;HST;U	V ++H57JJ1 K 
 
 	
 
 ++8O8O7P+Q++H57JJ1 K 
 
 	
 
)	X 	X
	V 	Vs   D 0D D	Dc                $   |j                  t        j                  g      }|j                  |dd       t	        j
                  t	        j                  |            }|j                  |      }|j                  |      }||cxk(  rddik(  sJ  J y )Nr   r\   pickle)r%   r   r&   rs   r  loadsdumpsrt   )r+   r   r,   restored_storagestorage_system_attrsrestored_storage_system_attrss         r   test_pickle_storagez#StorageTestCase.test_pickle_storage  s    ++8O8O7P+Q%%hx@!<<W(=>&==hG(8(O(OPX(Y%#'DYPXHYYYYYYr   c                   |j                  t        j                  g      }|j                  |      }|j	                  |t
        j                         |j	                  |t
        j                         t        j                  t              5  |j	                  |t
        j                         d d d        t        j                  t              5  |j	                  |t
        j                         d d d        t        j                  t              5  |j	                  |t
        j                         d d d        y # 1 sw Y   xY w# 1 sw Y   XxY w# 1 sw Y   y xY wr1   )r%   r   r&   rC   check_trial_is_updatabler   r   r   r;   r<   r
   r  ru  r   r   s       r   test_check_trial_is_updatablez-StorageTestCase.test_check_trial_is_updatable  s   ++8O8O7P+Q++H5((:3E3EF((:3E3EF]]34 	H,,XzG	H ]]34 	J,,Xz7H7HI	J ]]34 	L,,Xz7J7JK	L 	L	H 	H	J 	J	L 	Ls$   !D:!E!E:EEENr   r   rV   rW   )r   r   r   r   r   r   rV   rW   )r   r   r_  z	list[str]rV   rW   )r   r   rQ  r   r   zlist[float]rV   rW   )7r9   
__module____qualname__r/   r5   r8   rG   rK   rM   rZ   rj   ro   rv   ry   r{   r   r;   markparametrizer   r   r   r   r   r   r   r  r  r#  r)  r,  r7  r?  rB  rG  rK  rN  rU  r[  r^  rf  ro  r{  r  r  r  r   rX   r&   r   r  r  r  r  r  r   r   r   r   r   !   st   _ Nb+(
$9&D6O0?Q0A>8Jt [["
(,,.,(,,.	)HT1a,@(4QRTUBV+WX4J"4J084JIQ4J	4J	4Jl;(ZT9D,=U~:%_N(ET?
=!DF>
?FB@S 5: [[]c3Z#s,DE<"<1:<	< F</#$J/4?:2 [[[>+B+BND[D[*\][[%,$%-%5\3U6]35\N6]O	

H"H/=HGRH	H
 ^H(-
2ZLr   r   MNISTr   imageclassificationr=  tags)rc   nonejson_serializablezdict[str, JSONSerializable]rg   r  r%  )	zeropir   znegative maxr  znegative minr  znegative infr%  c                    t        j                  |       rt        j                  |      S t        j                  |      ry| |k(  S )NF)mathr(  )r`  ra  s     r   is_equal_floatsr    s2    zz!}zz!}zz!}6Mr   c                    | j                  t        j                  g      }t        j	                         D ]8  \  }}| j                  |||       t        | j                  |      |   |      r8J  y r1   )r%   r   r&   FLOAT_ATTRSrh   r_   r  r`   r   r,   r\   r]   s       r   rm   rm     sp    ''N4K4K3L'MH "'') S
U##Hc59w;;HEcJERRRSr   c                    | j                  t        j                  g      }t        j	                         D ]8  \  }}| j                  |||       t        | j                  |      |   |      r8J  y r1   )r%   r   r&   r  rh   rs   r  rt   r  s       r   rx   rx     sp    ''N4K4K3L'MH "'') U
U%%hU;w==hGLeTTTUr   c                f   | j                  t        j                  g      }t        j	                         D ]y  }t        j                  |      r| j                  |      }| j                  |t        j                  |f       | j                  |      j                  }|J t        ||      ryJ  y )Nr   r   )r%   r   r&   r  r   r  r(  rC   r   r   r   r   r]   r  )r   r,   r]   r   	set_values        r   r   r     s    ''N4K4K3L'MH##% 1::e++H5&&xz7J7JTYS[&\%%h/55	$$$y%0001r   c                0   | j                  t        j                  g      }| j                  |      }t        j                         D ]  \  }}t        ||      }t        |f      }||fD ]  }t        |t              rt        j                  |      s)|j                  j                  |z   }|j                  |      }	| j                  |||	|       t        | j!                  ||      |	      sJ | j#                  |      j$                  |   |k(  rJ   y )Nr   r  r  )r%   r   r&   rC   r  rh   r	   r   
isinstancer  isfinite	__class__r9   r   r  r  r   r   r   )
r   r,   r   r\   r]   float_distributioncategorical_distributionrc  re  internal_reprs
             r   r  r    s   ''N4K4K3L'MH''1H!'') Y
U.5uE#:E8#L /1IJ 	YL,(9:4==QVCW%//883>J(99%@M##Hj-V"7#:#:8Z#PR_```$$X.<<ZHLXXX	YYr   c                ,   | j                  t        j                  g      }| j                  |      }t	        t
        j                               D ]B  \  }}| j                  |||       t        | j                  |      j                  |   |      rBJ  y r1   )r%   r   r&   rC   	enumerater  r   r&  r  r   r   )r   r,   r   r   r]   s        r   r+  r+    s    ''N4K4K3L'MH''1Hk0023 Z5,,Xq%@w00:NNqQSXYYYZr   c                   | j                  | j                  t        j                  g            }t        j                         D ]8  \  }}| j                  |||       t        | j                  |      |   |      r8J  y r1   )	rC   r%   r   r&   r  rh   r;  r  r/  r   r   r\   r]   s       r   rA  rA    s~    ''  ^-D-D,E FH
 "'') S
U##Hc59w;;HEcJERRRSr   c                   | j                  | j                  t        j                  g            }t        j                         D ]8  \  }}| j                  |||       t        | j                  |      |   |      r8J  y r1   )	rC   r%   r   r&   r  rh   rJ  r  rE  r  s       r   rM  rM    s~    ''  ^-D-D,E FH
 "'') U
U%%hU;w==hGLeTTTUr   c                   t        j                  |      }i }i }t        |      D ]  }d| }	|/|j                  t        j
                  t        j                  g      }| j                  |f|	      }
| j                  |
d|       | j                  |
d|       i }t        |      D ]:  }t        |      }||_        | j                  |
|      |_        |||j                  <   < |||
<   t        |	|d|id|i|
      ||
<    ||fS )Nztest-study-name-r7   ur#   )r    rQ  r   r   r,   )rv  rw  r   choicer   r&   rX   r%   r_   rs   rx  r   rC   r   r   )r   r   r   r   rQ  ry  study_id_to_frozen_studystudy_id_to_trialsr   r    r,   rF   jr1  s                 r   r   r     s-    d#I79<>7^ 
's+
!((.*A*A>CZCZ)[\I++	|PZ+[##Hc15%%hQ7w 	,A#I.EEL%66xGEO&+F5??#		,
 (.8$-8!Qxq.
 *
* $%777r   c                   | j                  dd      t        dd      f| j                  dd      t        ddd      f| j                  g d      t        d      f| j                  dd      t        dd      f| j                  dd	g      t        d
      fd}ddd dddgdd}| j                  t              }i }i }i }i }i }|j                         D ]&  \  }	\  }
}| j                  ddg      s|
||	<   |||	<   ( |j                         D ]<  \  }	}
| j                  ddg      r|
|d|	z   <   | j                  ddg      s5|
|d|	z   <   > t        | j                  dd            D ]+  }| j                  ddg      s| j                  dd      ||<   - t        d||t        j                  k(  r| j                  dd      nd t        j                         |j                         rt        j                         nd |||||d      S )Nr   r   r   T)r
  )CatACatBCatCr  r  )r  r  )r   paramBparamCparamDparamEvalueAr   r  r  r  )attrAattrBattrCattrDFusr_sys_   r   i)r   r   r]   r   r   r   r   r   r   r   r   )uniformr	   r  r   r   rh   r   randintr   r   r   r   r   r   )ry  example_paramsexample_attrsr   r   r   r   r   r   r\   r]   distr   s                r   rx  rx  /  sB   $$Q*,=a,CD$$Q*,=a,MN56#$<=
 $$R+->r1-EF##S#J/1H1TU	N $)G=M3NO	M Z(EFMJ#%L,224 &]eTT5M*F3K!%M#& $))+ /
UT5M*',Jv|$T5M*).L#&	/
 9$$Q+, @T5M*%.%6%6sB%?"@ ,1Z5H5H,HiR(d||~,1,=,=,?(,,.T#!/ r   )rV   r   )r`  r   ra  r   rV   boolr  r   )r   r   r   r   r   r   r   r   rQ  zStudyDirection | NonerV   z@tuple[dict[int, FrozenStudy], dict[int, dict[int, FrozenTrial]]])ry  zrandom.RandomrV   r   )9
__future__r   rk  r   r  r  rv  systimer   typingr   numpyr'  r;   r=   optuna._typingr   optuna.distributionsr   r	   optuna.exceptionsr
   optuna.storagesr   optuna.storages._baser   optuna.study._frozenr   optuna.study._study_directionr   optuna.trialr   r   fixturer   r   rR   r   rg   __annotations__r  
float_infor   r  r   r  r  rm   rx   r   r  r+  rA  rM  r   rx  r   r   r   <module>r     sz   "      
      + 8 2 6 ' ; , 8 $ #  ~L ~LB *
 ,1GEU;VW.*  
''>>^^'''>>^^'''<5\M<
SU	1Y$ZSU  (,888 8 	8
 %8 F8D/r   