U
    *if;                     @   s@  d dl Zd dlZd dlmZmZmZmZmZ d dl	m
Z d dlmZmZ d dlmZmZmZ d dlmZ dd Zdd	 Zd
d Zdd Zejdeegejddedddddddddejg
i fdejdddddddddejg
i fdejdddddddddejg
i fdejdddddd ddd!ejg
d"d#ifd$ej dddd%d&d'ddd(ejg
d"d#ifd)ej!ddddddddd*ejg
i fgej"d+d,d- Z#ejdeegd.d/ Z$ejd0d1dddd2dd3ddejejg
fd4dddd5d6d7ddejejg
fgd8d9 Z%ejd:d;dddddd<ddddg
gd=ddddd>dd>dddg
ggd?d@ Z&dAdB Z'ejdCd#dd#d#dDdEgfdFdd#d dDd#gfgejdGd)g fdHd(gfgdIdJ Z(ejdKeeegejdLd#dEdMgejdNd#d#gd d#gdOd#dEgd d#gdOd#gdP ejd#dEejge)e*dQdP dOgdRdS Z+ejdTd#d#d gd#gfd#dEd d#gd#dEgfdEd#d gd#gfdEdEd d#gdEdEgfdUdMe*dMe)e*dUdMdMgdU  fdMdUe*dUdUgdU fd d e,g e,g fd#d e,g e,g fd d#d gd gfg	dVdW Z-ejdXed d#dEgd#dEdEgdYdEed dZdgej.d#dEdEge*dDgd[dgd\d]ej/d^feejd#dEejge)e*dQd_ d#gd` dEgda  e*d_dbdMedcdcddd>d>dedd<dfdgdhdidjdkdldmdndogej.d#gd` dEgda  e*d_gd[dgd\d]ej/d^fgdpdq Z0ejdrdsdtdudvdwdxdydzd{g	ejdLd#dEdDdQdUd|d}gd~d Z1dd Z2dd Z3dS )    N)	DataFrame
MultiIndexSeriesconcat
date_range)BaseIndexerFixedForwardWindowIndexer)ExpandingIndexerFixedWindowIndexerVariableOffsetWindowIndexer)BusinessDayc               	   C   sF   G dd dt } |  }tjtdd ttd| W 5 Q R X d S )Nc                   @   s   e Zd Zdd ZdS )z8test_bad_get_window_bounds_signature.<locals>.BadIndexerc                 S   s   d S N )selfr   r   X/home/mars/bis/venv/lib/python3.8/site-packages/pandas/tests/window/test_base_indexer.pyget_window_bounds   s    zJtest_bad_get_window_bounds_signature.<locals>.BadIndexer.get_window_boundsN__name__
__module____qualname__r   r   r   r   r   
BadIndexer   s   r   zBadIndexer does not implementmatch   )r   pytestraises
ValueErrorr   rangerolling)r   indexerr   r   r   $test_bad_get_window_bounds_signature   s    r    c                  C   s<   t td} t }| | }|   }t|| d S )N
   )r   r   r	   r   meanZ	expandingtmassert_series_equal)sr   resultexpectedr   r   r   test_expanding_indexer#   s
    r(   c                  C   sn   dddddg} t dtdi}G dd dt}|d| d}|| }t dd	d
dddgi}t|| d S )NTFvaluesr   c                   @   s   e Zd Zdd ZdS )z3test_indexer_constructor_arg.<locals>.CustomIndexerc                 S   sl   t j|t jd}t j|t jd}t|D ]:}| j| rLd||< |d ||< q(|||< || j ||< q(||fS Ndtyper      )npemptyint64r   use_expandingwindow_sizer   
num_valuesmin_periodscenterclosedstartendir   r   r   r   1   s    
zEtest_indexer_constructor_arg.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   CustomIndexer0   s   r;   r-   r2   r1                 ?      @      $@r   r   r   r   sumr#   assert_frame_equal)r1   dfr;   r   r&   r'   r   r   r   test_indexer_constructor_arg+   s    rE   c                  C   sf   t dtdi} G dd dt}|dd}| j|dddd	 }t dd
ddddgi}t|| d S )Nr)   r   c                   @   s   e Zd Zdd ZdS )z8test_indexer_accepts_rolling_args.<locals>.CustomIndexerc                 S   sz   t j|t jd}t j|t jd}t|D ]H}|rZ|dkrZ|dkrZ|dkrZd||< |||< q(|||< || j ||< q(||fS )Nr+   r-   both   r   )r.   r/   r0   r   r2   r3   r   r   r   r   G   s    
zJtest_indexer_accepts_rolling_args.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r;   F   s   r;   r-   r2   TrF   )r6   r5   r7   r=   r>   r@   r?         @rA   )rD   r;   r   r&   r'   r   r   r   !test_indexer_accepts_rolling_argsC   s    
rJ   constructorzfunc,np_func,expected,np_kwargscountr?   g       @minr=   r>   rI   g      @g      @g       @max      Y@g      "@stdgK@gDmK@g!?J@g'e?Zddofr-   varg[{@@g[{@g[{Ħ@g      ?mediang      !@z ignore:min_periods:FutureWarningc              	      s@  t d}d|d< tdd}d}tjt|d$ | |j|dd	}t||  W 5 Q R X d
}tjt|d$ | |j|dd}t||  W 5 Q R X | |j|dd}t|| }	| |}t	|	| | |
 fdd}
t	|	|
 |dkrdnd }| |j||d}t|| }| |
 fdd}t	|| d S )Nr@   rO   r      rH   z.Forward-looking windows can't have center=Truer   T)windowr6   zAForward-looking windows don't support setting the closed argumentright)rT   r7   rG   rT   r5   c                    s    | fS r   r   xnp_func	np_kwargsr   r   <lambda>       z-test_rolling_forward_window.<locals>.<lambda>rL   r   c                    s    | fS r   r   rW   rY   r   r   r\      r]   )r.   aranger   r   r   r   r   getattrr#   assert_equalapply)rK   funcrZ   r'   r[   r)   r   r   r   r&   Z	expected2r5   Zrolling3Zresult3Z	expected3r   rY   r   test_rolling_forward_windowY   s,    9

rc   c                 C   sf   t d}d|d< tdd}| |j|dd}| }| ddd	d
ddddt jt jg
}t|| d S )Nr@   rO   r   rH   rS   rV   r=   go@g\QJ@g@gv-@gA@)r.   r^   r   r   Zskewnanr#   r`   )rK   r)   r   r   r&   r'   r   r   r   test_rolling_forward_skewness   s&    

re   zfunc,expectedcovg     @X@g     @Wcorrg?gh"?g1=c                 C   s   t ddd}|d }d|d< t j||gdd}tdd	}t|j|dd
}t||  jt	d dfdf }|j
dd}t|}|j|_t|| d S )Nr!   r-   rG   d   )r   r   )ZaxisrS   rH   rV   r   T)Zdrop)r.   r^   ZreshapeZconcatenater   r   r   r_   locsliceZreset_indexr   namer#   r`   )rb   r'   Zvalues1Zvalues2r)   r   r   r&   r   r   r   test_rolling_forward_cov_corr   s    
rm   zclosed,expected_datarU   g      (@leftg      @c                 C   s^   t ddd}ttd|d}td}t||d}|j|| d }t||d}t|| d S )NZ2020r!   )Zperiods)indexr-   )ro   offset)r7   )	r   r   r   r   r   r   rB   r#   rC   )r7   Zexpected_dataro   rD   rp   r   r&   r'   r   r   r   &test_non_fixed_variable_window_indexer   s    rq   c                  C   sT   t dd d d dgi} tdd}| j|dd }t dddddgi}t|| d S )	Nb   rG   rH   r   rV   r=   r>   )r   r   r   rL   r#   rC   )rD   r   r&   r'   r   r   r    test_fixed_forward_indexer_count  s
    
rt   )	end_valuer)   rS   rG   rh   )rb   argsZquantilec           
         sn   G  fdddt }dddddg}tdtdi}|d|d}t|||| }td|i}	t||	 d S )	Nc                       s   e Zd Z fddZdS )z0test_indexer_quantile_sum.<locals>.CustomIndexerc                    sr   t j|t jd}t j|t jd}t|D ]@}| j| rRd||< t|  d||< q(|||< || j ||< q(||fS r*   )r.   r/   r0   r   r1   rN   r2   r3   ru   r   r   r     s    
zBtest_indexer_quantile_sum.<locals>.CustomIndexer.get_window_boundsNr   r   rw   r   r   r;     s   r;   TFr)   r   r-   r<   )r   r   r   r_   r   r#   rC   )
ru   r)   rb   rv   r;   r1   rD   r   r&   r'   r   rw   r   test_indexer_quantile_sum  s    rx   indexer_classr2      df_data)arr         c                 C   sX   t |}d}| |d}|j}t|D ].}|dd j|dd  |j|ks$tq$d S )NrS   rH   r|   rr   r-   rV   )r   r2   r   groupbyr   r"   AssertionError)ry   r2   r{   rD   Z
num_trialsr   Zoriginal_window_sizer:   r   r   r   0test_indexers_are_reusable_after_groupby_rolling-  s    
r   z5window_size, num_values, expected_start, expected_endr   c                 C   s^   t | d}|j|d\}}tj|t|dd tj|t|dd t|t|ksZtd S )NrH   )r4   F)Zcheck_dtype)r   r   r#   Zassert_numpy_array_equalr.   arraylenr   )r2   r4   Zexpected_startZexpected_endr   r8   r9   r   r   r   !test_fixed_forward_indexer_boundsF  s
    
r   zdf, window_size, expected)rr   r|   g      ?r|   )namesrr   )ro   rl   r,      rs      )rr   r|   cg@g      @g      @g      )@g      *@g      +@g      ,@g      -@g      .@g      /@g      0@g     0@g      1@c                 C   s6   t |d}| dd j|dd }t|| d S )NrH   r|   rr   r-   rV   )r   r   r   r"   r#   r$   )rD   r2   r'   r   r&   r   r   r   0test_rolling_groupby_with_fixed_forward_specific`  s    6
r   
group_keys)r-   )r-   rG   )rG   r-   )r-   r-   rG   )r-   rG   r-   )r-   r-   rG   rG   )r-   rG   rS   rG   rS   )r-   r-   rG   r-   r-   rG   r-   r-   rG   r-   r-   rG   )r-   rG   rS   r-   rG   rS   r-   rG   rS   r-   rG   rS   r-   rG   rS         c                    s   t tt| tjt| tjdd tjt| tjdd}t d}|	dd j
|dd }dd	g|j_|	dddd	g }t fd
d|D }|dd	gd }t|| d S )Nr+      )r|   rr   r   rH   r|   rr   r-   rV   r   c                    s4   g | ],\}  j  fd dtt D dqS )c                    s*   g | ]"} d  j ||  jddqS )rr   r-   )Z	min_count)ZilocrB   ).0r:   )gr2   r   r   
<listcomp>  s   zKtest_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>.<listcomp>)rr   )Zassignr   r   )r   _rH   )r   r   r     s   
z@test_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>)r   r.   r   listr^   r   float64r0   r   r   r   rB   ro   r   r   Z	set_indexr#   r$   )r   r2   rD   r   r&   groupsZmanualr   rH   r   ,test_rolling_groupby_with_fixed_forward_many  s"    

r   c               	   C   s   G dd dt } |  }td|}d}tjt|d |  W 5 Q R X tjt|d tt| W 5 Q R X tjt|d |j	dd W 5 Q R X tjt|d |j
dd W 5 Q R X d S )Nc                   @   s   e Zd Zdd ZdS )z4test_unequal_start_end_bounds.<locals>.CustomIndexerc                 S   s   t dgt ddgfS Nr-   rG   r.   r   r   r4   r5   r6   r7   r   r   r   r     s    zFtest_unequal_start_end_bounds.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r;     s   r;   r-   r8   r   Tpairwiser   r   r   r   r   r   r"   nextiterrg   rf   r;   r   Zrollr   r   r   r   test_unequal_start_end_bounds  s    r   c               	   C   s   G dd dt } |  }tddg|}d}tjt|d |  W 5 Q R X tjt|d tt| W 5 Q R X tjt|d |j	dd W 5 Q R X tjt|d |j
dd W 5 Q R X d S )Nc                   @   s   e Zd Zdd ZdS )z4test_unequal_bounds_to_object.<locals>.CustomIndexerc                 S   s   t dgt dgfS r   r   r   r   r   r   r     s    zFtest_unequal_bounds_to_object.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r;     s   r;   r-   zstart and endr   Tr   r   r   r   r   r   test_unequal_bounds_to_object  s    r   )4numpyr.   r   Zpandasr   r   r   r   r   Zpandas._testingZ_testingr#   Zpandas.api.indexersr   r   Zpandas.core.indexers.objectsr	   r
   r   Zpandas.tseries.offsetsr   r    r(   rE   rJ   markZparametrizer   rd   rM   rN   rP   rQ   rR   filterwarningsrc   re   rm   rq   rt   rx   r   r   r   r   r   Zfrom_arraysr   r   r   r   r   r   r   r   r   <module>   sv  
 "5
(



	 " ("
 4
