# 将数据按照一定的形式进行分割，分割出来的每个数据文件进行词频统计
# 这里展示的是按照 地区*时间 颗粒度组合下的分割方式。
# 将指定地区的微博用户按照 “用户/地区” * “年/季度/月/天” 的形式进行分割保存
# 数据文件的格式：UID, location, 微博ID, 发表时间戳， 2010-1-1开始的月数， 2010-1-1开始的周数， 微博文本

import os
import sys
import datetime, time
import pandas as pd
from enum import Enum


# 按地点（省/市）和用户
class LocGranularity(Enum):
    PROVINCE = "p"
    CITY = "c"
    UID = "u"

# 不同时间颗粒度，年/季度/月/天
class TimeGranularity(Enum):
    DAILY = "d"
    MONTHLY = "m"
    QUARTERLY = "q"
    YEARLY = "y"


# ===================== 参数设置 =====================
inPath = "../data/dalian/"

#每个地区是一个目录，目录下的每个文件是一天的微博数据
outPath = "../data/timeloc/dalian/"

# 数据输出以 "用户/地区 * 不同时间颗粒度" 的形式输出到文件

# 是以用户UID为单位, 还是以整个地区(地市级)为单位
locType =  LocGranularity.CITY

# 时间颗粒度的类型，年/季度/月/天
twType = TimeGranularity.MONTHLY

# ===================== 参数设置 =====================


#将每年每个地区的数据一次性写入文件，并清空该词典的数据内容
def saveBuf(dicData):
    if dicData:
        for kitem in dicData.keys():
            with open(kitem, 'a', encoding='utf-8') as outf:
                outf.write(dicData[kitem])
                outf.flush()
        dicData.clear()
    return


if(not os.path.exists(inPath)):
    print("DIR ", inPath, " 不存在!")
    sys.exit(0)
    
buf_data = {}

#读入每个文件，按照微博用户ID 和 年月 的形式进行保存
for fdata in os.listdir(inPath):
    st = time.time()
    infile = inPath + fdata 
    print(infile)
    
    #从数据文件中读入，解析每行内容
    with open(infile, encoding='utf-8') as freader:
         num_lines = 0
         for fline in freader:         
             line_content = fline.rstrip('\r\n')
             seg = line_content.split('\t')
             if(len(seg) != 7):
                 continue
            
             #获取微博用户UID和注册地
             uid = seg[0]
             loc = seg[1]
             
             # 如果是设定了 省市两级的话，需要有地市级名称
             if( (loc.find(" ") < 0) and (locType == LocGranularity.CITY)):
                 continue
             
             #获取省市信息
             province = ''
             city = ''
             if( loc.find(" ") < 0 ):
                 province= loc.strip()
             else:
                 locseg = loc.split(' ')
                 province = locseg[0]
                 city = locseg[1]                
             
             if((province == '海外') or (province == '其他') or (province == 'noop')):
                 continue
                
             # 按照不同时间颗粒度，转换为文件名
             dateArray = time.localtime(int(seg[3]))
             dstring = time.strftime("%Y-%m-%d %H:%M:%S", dateArray)
             
             if(twType == TimeGranularity.YEARLY):
                 strTW = str(dateArray.tm_year)
             elif(twType == TimeGranularity.QUARTERLY):
                 strTW = str(dateArray.tm_year)+'Q'+str(pd.Timestamp(dstring).quarter)
             elif(twType == TimeGranularity.MONTHLY):
                 strTW = time.strftime("%Y-%m", dateArray)
             else:
                 strTW = time.strftime("%Y-%m-%d", dateArray)
                
             # dateArray = datetime.datetime.utcfromtimestamp(int(seg[3]))
             # strTW = dateArray.strftime('%Y-%m')
                
             # 获得微博文本
             text = seg[6]
            
             # 以用户ID/地区(地市级) 和 不同的时间颗粒度 组合成输出文件
             if(locType == LocGranularity.PROVINCE):
                 outfile = outPath + province + '_' + strTW + ".txt"
             elif(locType == LocGranularity.CITY):
                 outfile = outPath + loc + '_' + strTW + ".txt"
             else:
                 outfile = outPath + uid + '_' + strTW + ".txt"
                 
                 
             # 只输出微博文本
             buf_data[outfile] = buf_data.get(outfile, '') + text + '\n'
    
             num_lines += 1
            
             # 达到一定数量后，一次性写入文件
             if(num_lines >= 100000):
                 print('.', end='')
                 saveBuf(buf_data)
                 num_lines = 0
                 
    print()
    saveBuf(buf_data)
    
    et = time.time()
    print(fdata,":",(et-st),"s")


print('ALL DONE!')

