<gh_stars>0
import os
import re
import gzip
import pandas as pd
import logging
from datetime import datetime
from dateutil import tz

logger = logging.getLogger(__name__)

FILE_REGEX = re.compile(r'(?P<year>[0-9]{4})-(?P<month>[0-9]{2})-(?P<day>[0-9]{2})\s(?P<hour>[0-9]{2})_(?P<minute>[0-9]{2})_(?P<second>[0-9]{2})(?P<extension>\..*)')
CHUNKSIZE = 10 ** 6

def process(study, subject, read_dir, date_from, output_tz, input_tz):
    # Instantiate an empty dataframe
    df = pd.DataFrame.from_records([])

    for root_dir, dirs, files in os.walk(read_dir):
        files[:] = [ f for f in files if not f[0] == '.' ]
        dirs[:] = [ d for d in dirs if not d[0] == '.' ]

        for file_name in sorted(files):
            file_name, extension = verify(file_name)
            if file_name is not None:
                file_path = os.path.join(root_dir, file_name)
                for data in get_data(file_path, extension):
                    data_list = parse(data, date_from, output_tz,
                            input_tz, file_path, file_name)
                    df = df.append(data_list, ignore_index=True, sort=False)

    return df

def insert_cols(df):
    cols = df.columns.tolist()

    if 'Incoming Call' not in cols:
        df.insert(0, 'Incoming Call', 0)
    if 'Outgoing Call' not in cols:
        df.insert(1, 'Outgoing Call', 0)
    if 'Missed Call' not in cols:
        df.insert(2, 'Missed Call', 0)

    return df

def process_seconds(df):
    df.index.name = None
    dfg = df.groupby(['day', 'weekday', 'timeofday', 'UTC_offset', 'call type'])

    df = dfg.size().reset_index(name='counts')
    df = df.pivot_table(index=['day','weekday',
        'timeofday','UTC_offset'], columns='call type', values='counts').fillna(0)

    df = insert_cols(df)

    df = df.rename(columns={
        "Incoming Call": "incoming_call_counts",
        "Outgoing Call": "outgoing_call_counts",
        "Missed Call": "missed_call_counts"
    })

    df['incoming_call_counts'] = df['incoming_call_counts'].astype(int).astype(str)
    df['outgoing_call_counts'] = df['outgoing_call_counts'].astype(int).astype(str)
    df['missed_call_counts'] = df['missed_call_counts'].astype(int).astype(str)

    df_2 = dfg['duration in seconds'].agg(['min','max','sum','mean']).reset_index()
    df_min = df_2.pivot_table(index=['day','weekday'],columns='call type', values='min').fillna(0)
    df_min = insert_cols(df_min)
    df['incoming_call_duration_in_seconds_min'] = df_min['Incoming Call']
    df['outgoing_call_duration_in_seconds_min'] = df_min['Outgoing Call']

    df_max = df_2.pivot_table(index=['day','weekday'],columns='call type', values='max').fillna(0)
    df_max = insert_cols(df_max)
    df['incoming_call_duration_in_seconds_max'] = df_max['Incoming Call']
    df['outgoing_call_duration_in_seconds_max'] = df_max['Outgoing Call']

    df_sum = df_2.pivot_table(index=['day','weekday'],columns='call type', values='sum').fillna(0)
    df_sum = insert_cols(df_sum)
    df['incoming_call_duration_in_seconds_sum'] = df_sum['Incoming Call']
    df['outgoing_call_duration_in_seconds_sum'] = df_sum['Outgoing Call']

    df_mean = df_2.pivot_table(index=['day','weekday'],columns='call type', values='mean').fillna(0)
    df_mean = insert_cols(df_mean)
    df['incoming_call_duration_in_seconds_mean'] = df_mean['Incoming Call']
    df['outgoing_call_duration_in_seconds_mean'] = df_mean['Outgoing Call']

    df_3 = dfg['hashed phone number'].nunique().reset_index(name='unique')
    df_unique = df_3.pivot_table(index=['day','weekday'],columns='call type', values='unique').fillna(0)
    df_unique = insert_cols(df_unique)
    df['incoming_call_unique_numbers'] = df_unique['Incoming Call'].astype(int).astype(str)
    df['missed_call_unique_numbers'] = df_unique['Missed Call'].astype(int).astype(str)
    df['outgoing_call_unique_numbers'] = df_unique['Outgoing Call'].astype(int).astype(str)

    df.columns.name = None

    return df.round(3)

# Verify the file based on its filename
def verify(file_name):
    match = FILE_REGEX.match(file_name)
    if match:
        return file_name, match.group('extension')
    else:
        return None, None

# Unzip the file and parse as csv
def gz_to_df(file_path):
    try:
        with gzip.open(file_path, 'rb') as f:
            if hasattr(f, 'read'):
                return csv_to_df(f)
            else:
                return None
    except Exception as e:
        logger.error(e)
        return

# Convert a comma-separated-string object to pandas csv dataframe
def csv_to_df(file_path):
    try:
        for chunk in pd.read_csv(file_path, keep_default_na=False, engine='c',
                skipinitialspace=True, error_bad_lines=False, chunksize=CHUNKSIZE,index_col=False):
            yield chunk
    except Exception as e:
        logger.error(e)
        return

# Check the file extension, and act accordingly
def get_data(file_path, extension):
    if extension == '.csv':
        return csv_to_df(file_path)
    elif extension == '.csv.gz':
        return gz_to_df(file_path)
    else:
        logger.error('Incompatible file extension %s' % extension)
        return

# Return timestamp in the timezone
def process_datetime(row_timestamp, input_tz, output_tz):
    timestamp = pd.Timestamp(row_timestamp, tz = input_tz)
    return timestamp.tz_convert(output_tz)

# Parse and process data
def parse(df, date_from, output_tz, input_tz, file_path, filename):
    if df is None or len(df) == 0:
        logger.error('Could not open file %s' % file_path)
        return

    df['$date_to'] = df['UTC time'].apply(lambda x: process_datetime(x, input_tz, output_tz))
    return df
