
    tf>                     8   d dl mZ d dlZ	 d dlZd dl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Zd dlZd dlZd dlZd dlZd dlZd dlZd dlmZmZ d Zd Z G d de      Z G d d	e      Z G d
 de      Z  G d de      Z!	 ddZ"ddZ#d Z$y# e$ r d dlZY w xY w)    )print_functionN)gmtimestrftimec           	          t        dj                  t        j                  | ||                   t	        j
                  d       y )N    )printjoin	tracebackformat_exceptionos_exit)exc_type	exc_valueexc_tracebacks      \/var/www/html/software/conda/envs/higlass/lib/python3.12/site-packages/clodius/save_tiles.pyhandle_exceptionr      s-    	"''),,Xy-P
QRHHQK    c                    t        j                   t         j                  t         j                         | j                         r|j                  sH	 | j                  d      \  }}}|j                  |||       | j                         s;|j                  sH|j                          y # t        t        f$ r t        d       Y -t        j                  $ r |j                          Y ow xY w)Nr   )timeoutz
Exiting...)signalSIGINTSIG_IGNemptyvaluegetsave_binned_tileKeyboardInterrupt
SystemExitr	   queueEmptyflush)q
tile_saverfinished
zoom_leveltile_pos	tile_binss         r   tile_saver_workerr)       s    
MM&--0ggi	01a0@-Z9''
HiH ggi  ":. 	,{{ 		s   )B( (C)"C)(C)c                   B    e Zd Z	 	 d
dZd Zd Zd Zd Zd Zd Z	d Z
y	)	TileSaverc                 J    || _         || _        || _        || _        || _        y N)max_data_in_sparsebins_per_dimensionnum_dimensionsprint_statusinitial_value)selfr.   r/   r0   r1   r2   s         r   __init__zTileSaver.__init__4   s1     #5 #5,(*r   c                      y r-    )r3   tiles     r   	save_tilezTileSaver.save_tileF       r   c           	          dj                  |dj                  t        t        |                  }||d}| j	                  |       y )Nz{}.{}.tile_id
tile_value)formatr
   mapstrr8   )r3   r&   tile_position	tile_datar=   r7   s         r   make_and_save_tilezTileSaver.make_and_save_tileI   s<     ..SXXc#}6M-NO #)<tr   c                    | j                   g| j                  | j                  z  z  }|j                         D ]J  \  }}t	        t        t        |      D 	
cg c]  \  }	}
|
| j                  |	z  z   c}
}	            }|||<   L t        | j                         dk(  r3| j                  |||D cg c]  }t        |d   d       c}||d       y | j                  ||t        t        j                  j                  |D cg c]  }|D cg c]  }t        |d       c} c}}            ||d       y c c}
}	w c c}w c c}w c c}}w )Nr   r      dense	min_value	max_value)r2   r/   r0   itemsintsum	enumeratelenrD   roundlistitchainfrom_iterable)r3   r&   rB   r(   rI   rJ   initial_valuesbin_posvalibpindexvys                 r   save_dense_tilezTileSaver.save_dense_tileT   sE    ,,-##t':'::
 (oo/ 	(NWcIgDVW51bR$11Q66WXE %(N5!		( t!!"a'##6DEeAaD!nE!*!* ##!..?MN!15aeAqk5N
 "+!* X F 6Ns$   D6.D< 	E	EEEc                    g }|j                         D ]^  \  }}t        |      dk(  r#|t        t        t        |            |d   ggz  }7|t        t        t        |            t        |      ggz  }` | j                  |||||d       y )Nr   r   )sparserI   rJ   )rK   rO   rQ   r@   floatrD   )	r3   r&   rB   r(   rI   rJ   shownrV   bin_vals	            r   save_sparse_tilezTileSaver.save_sparse_tilez   s     "+//"3 	FWg7|q 4E7 34gajABB4E7 34d7mDEE		F 	99M	
r   c                     t        |      }t        |      }| j                  |||D cg c]  }t        |d       c}||d       yc c}w )aQ  
        Save a tile that has all of its data in one long array

        :param zoom_level: An integer zoom_level (0 for zoomed all the way out)
        :param tile_position: An n-dimensional array, where n is the number of dimensions
                              in the dataset.
        :param tile_data: The data in the tile.
        rF   rG   N)minmaxrD   rP   )r3   r&   rB   rC   rI   rJ   r[   s          r   save_tile_arrayzTileSaver.save_tile_array   sN     	N		N	/89!%1+9&&	
 :s   Ac           
         t        t        j                  t        j                  t        |j	                                     d            }t        t        j
                  t        j                  t        |j	                                     d            }t        |      | j                  k  r| j                  |||||       y | j                  |||||       y )Nr   )axis)rI   rJ   )
rQ   nprf   arrayvaluesre   rO   r.   rc   r]   )r3   r&   rB   r(   rJ   rI   s         r   r   zTileSaver.save_binned_tile   s    i.>.>.@)A BKL	i.>.>.@)A BKL	y>D333!!## "    ## ! r   c                       y r-   r6   r6   r   r   r"   zTileSaver.flush   r9   r   N)Fg        )__name__
__module____qualname__r4   r8   rD   r]   rc   rg   r   r"   r6   r   r   r+   r+   3   s3     $	$L
 
,*r   r+   c                        e Zd Z fdZ xZS )EmptyTileSaverc                 0    t         t        |   |||       y r-   )superrr   r4   )r3   r.   r/   r0   	__class__s       r   r4   zEmptyTileSaver.__init__   s    nd, 2N	
r   )rn   ro   rp   r4   __classcell__ru   s   @r   rr   rr      s    
 
r   rr   c                   *     e Zd Z fdZd Zd Z xZS )ColumnFileTileSaverc                     t         t        |   |||||       || _        t	        j
                         | _        d| _        || _        y Nr   )	rt   ry   r4   	file_pathcsioStringIObulk_txtbulk_txt_lenlog_file)	r3   r.   r/   r0   r|   r   r1   r2   ru   s	           r   r4   zColumnFileTileSaver.__init__   sI     	!41	
 # r   c           	      @   	 |d   dk(  r.| j                   j                  |d   dz   dz   dz   dz   dz          n|d   j                  d      }| j                   j                  t        t	        |d         dz         dz   t        t	        |d         dz         z   dz   t        t	        |d         dz         z   dz          | j                   j                  t        j                  |      dz          | j                   j                         }|d	kD  r| j                          y
y
)a]  
        if ('dense' in val['tile_value']):
            value_pos = col.defaultdict(list)
            dense_values = val['tile_value']['dense']
            dense_values = [(x,len(list(y))) for (x,y) in it.groupby(dense_values)]
            dense_values = [item for sublist in dense_values for item in sublist]
            val['tile_value']['dense'] = dense_values
            for i,value in enumerate(dense_values):
                value_pos[value] += [i]
            for key in value_pos:
                sorted_value_pos = sorted(value_pos[key])
                diffs = []
                diffs += [sorted_value_pos[0]]
                for i in range(len(sorted_value_pos)-1):
                    diffs += [sorted_value_pos[i+1] - sorted_value_pos[i]]

                value_pos[key] = diffs
            val['tile_value']['dense'] = value_pos.items()
        r=   tileset_info	1r;   r   r   
i N)	r   writesplitrA   rL   jsondumpstellr"   )r3   rW   ticurr_poss       r   r8   zColumnFileTileSaver.save_tile   s   *	& y>^+MMI 5 ;d BS H4 OPY%%c*BMMC1JN#c"Q%j1n%&  c"Q%j1n%	&
  	DJJsOd23==%%' gJJL r   c           	      `   | j                   j                         dkD  rI	 t        | j                  d      5 }|j	                  | j                   j                                d d d        d| _
        | j                   j                          t        j                         | _         y # 1 sw Y   DxY w# t        $ rv}| j                  `t        | j                  d      5 }|j	                  t        dt                            |j	                  |       d d d        n# 1 sw Y   nxY wY d }~d }~ww xY w)Nr   a%Y-%m-%d %H:%M:%S)r   r   openr|   r   getvalue	Exceptionr   r   r   r   closer}   r~   )r3   column_fileexfs       r   r"   zColumnFileTileSaver.flush  s    ==!#$$..#. @+%%dmm&<&<&>?@ @ @ $==,dmmS1 $Q)<fh GH$ $ $$sF   B. *B"B. "B+'B. .	D-7"D(5D	D(D 	D((D-)rn   ro   rp   r4   r8   r"   rv   rw   s   @r   ry   ry      s    !,:x(r   ry   c                   @     e Zd Z	 	 	 	 	 	 	 d fd	Zd Zd Zd Z xZS )ElasticSearchTileSaverc                     t         t        |   |||||       || _        t	        j
                         | _        d| _        || _        y r{   )	rt   r   r4   es_pathr}   r~   r   r   r   )	r3   r.   r/   r0   r   r   r1   r2   ru   s	           r   r4   zElasticSearchTileSaver.__init__$  sI     	$d4	
  r   c                    d|d   v r|d   d   }t        j                  t              }|D ]L  }t        | j                        dk7  r|t        |d         xx   |d   gz  cc<   9||d   xx   |d   gz  cc<   N g }|j                         D ]  \  }}t        |      }t        | j                        dk(  r|t        |      gz  }n|t        |      z  }|t        t        |            gz  }t        t        |d               D ]  }||D 	cg c]  }	|	|   	 c}	z  }  ||d   d<   | j                  |d   |       y c c}	w )Nr_   r>   r   r   r=   )coldefaultdictrQ   rO   r2   tuplerK   sortedr`   range
save_value)
r3   rW   sparse_values	value_possparse_valuevalue_xs_ysr   possrX   ps
             r   r8   z ElasticSearchTileSaver.save_tile:  s]   
 s<((-h7M-I - Dt))*a/eLO45,q/9JJ5l1o.<?2CC.	D K(0 8t d|t))*a/E%L>1K4;.Kc$i 011s47|, 8A$#7QAaD#77K88 +6Ch'I,	 $8s   D>
c                    | j                   j                  dj                  |             | j                   j                  t        j                  |      dz          	 | j                   j                         }|dkD  r| j                          yy)z
        if ('dense' in val['tile_value']):
            print val['tile_id'], len([x for x in val['tile_value']['dense'] if x > 0])
        z{{"index": {{"_id": "{}"}}}}
r   @KL N)r   r   r?   r   r   r   r"   )r3   doc_iddocr   s       r   r   z!ElasticSearchTileSaver.save_value\  sp     	<CCFKLDJJsOd23	$ ==%%'gJJL r   c           	      8   | j                   j                         dkD  r{	 t        d| j                  z   dz   | j                   j	                         | j
                         d| _        | j                   j                          t        j                         | _         y y # t        $ rv}| j                  `t        | j                  d      5 }|j                  t        dt                            |j                  |       d d d        n# 1 sw Y   nxY wY d }~d }~ww xY w)Nr   http://z/_bulkr   r   )r   r   save_to_elasticsearchr   r   r1   r   r   r   r   r   r   r   r   r}   r~   )r3   r   r   s      r   r"   zElasticSearchTileSaver.flush~  s    ==!#
$%,x7MM**,%% !"DMM! MMODM! $  $==,dmmS1 $Q)<fh GH$ $ $$s/   ?B 	D#"D5D:	DD	DD)NNNNNFN)rn   ro   rp   r4   r8   r   r"   rv   rw   s   @r   r   r   #  s0      !, -D D,r   r   c                    d}t        j                  ||      }t        j                  |d      }| D ]`  }|t        j                  dd|d   ii      dz   z  }|t        j                  |      dz   z  }t	        |      dkD  sOt        d|z   ||       d}b t        d	|       t	        |      d
kD  rt        d|z   ||       y y )Nr   _bulkrZ   _idr=   r   r   r   z	bulk_txt:r   )opr
   r   r   rO   r   r	   )	partitionelasticsearch_nodeselasticsearch_pathr1   r   es_urlput_urlrW   s           r   save_tile_to_elasticsearchr     s     HWW(*<=Fggfg&G DJJ%Y)@ABTIIDJJsOd**x=7"!)g"5xNH 
+x 
8}qi'18\J r   c           	      8   d}d}t        j                         }|s:	 t        j                  | |d      }|rt	        d||dt        |      |        d}|s9yy# t        $ r}|dz  }t	        d	|d
||t        j                         t        j                  |       |dkD  rrt        dd      5 }|j                  d| d       |j                  |       |j                          ddd       n# 1 sw Y   nxY wt	        dt        j                          Y d}~d}~ww xY w)a  
    Save some data to elastic search.

    The data should be a string suitable for bulk import by
    elasticsearch. The url should be the location of the index, document
    type, along with the _bulk destination.

    :param url: The elasticsearch url that will ingest the data
                e.g. localhost:9200/hg19/tiles/_bulk
    :param data: The data to import.
                e.g. {"index": {"_id": "blah", "my_json": {"x": 2}}}
    Fr      )datar   z
Savedz
len(data):T   z!
Error saving to elastic search (z), sleeping:fileiX  zunsaved.errr   zUNSAVED url:r   NzSlept too long, returning)slugidnicerequestspostr	   rO   r   sysstderrtimesleepr   r   r"   )	urlr   r1   savedto_sleepuidrr   r   s	            r   r   r     s    EH
++-C	ca8Aias4y#FE   	MH4ZZ JJx #~-- GGNC6GGDMGGI   1

C 	s0   5A 	DA
D)5C'	D'C0	,#DDc           	         | d   }| d   }t        j                  |dj                  t        t        |                  }t        j                  |      }t        j
                  |      s	 t        j                  |       ddj                  t        t        |            |di}|rIt        j                  |d	z   d
      5 }	|	j                  t        j                   |d             ddd       yt        |d
      5 }	|	j                  t        j                   |d             ddd       y# t        $ r&}t        d|t        j                         Y d}~d}~ww xY w# 1 sw Y   yxY w# 1 sw Y   yxY w)z
    Save a tile to a particular base directory.

    This function create the appropriate sub-directory based on the
    key.

    They key should be in the format (zoom_level, pos1, pos2...)
    e.g. (5,4,5)
    r   r   r;   zError making directories:r   N_sourcer<   z.gzwr   )indent)r   r
   r@   rA   dirnameexistsr   makedirsOSErrorr	   r   r   gzipr   r   r   r   )
r7   
output_dirgzip_outputkeyr>   outpathoutdiroeoutput_jsonr   s
             r   r8   r8     s6    q'CaJggj#((3sC="9:GZZ F99V	DKK 	sxxC6jQK YYw, 	7GGDJJ{156	7 	7 '3 	71GGDJJ{156	7 	7  	D -r

CC	D	7 	7	7 	7s0   )D) ='E9'E')	E2EEE$'E0)F)%
__future__r   collectionsr   	cStringIOr}   ImportErrorior   	itertoolsrR   r   numpyrj   r   os.pathpathr   r    r   r   r   r   r   r   r   r   r   r)   objectr+   rr   ry   r   r   r   r8   r6   r   r   <module>r      s    %      	      
   !
&C CL
Y 
`() `(Fl,Y l,` FKK(-` 7_  s   B 	BB