import sqlalchemy

import pandas as pd

from constants import variable_type,column_names,ALL

import logger

import numpy as np

import dask.dataframe as dd

def psql_connect(user, password, database='postgres', host='localhost', port=5432):

    '''Returns a connection and a metadata object'''

    # We connect with the help of the PostgreSQL URL

    # postgresql://federer:grandestslam@localhost:5432/tennis

    url = 'postgresql://{}:{}@{}:{}/{}'

    url = url.format(user, password, host, port, database)

    # The return value of create_engine() is our connection object

    connection = sqlalchemy.create_engine(url, client_encoding='utf8')

    return connection

def simple_sql_query(table,columns=['*'],where_condition=None):

    where = ''

    if where_condition is not None:

        where = ' WHERE {}'.format(where_condition)

    return 'SELECT {} FROM {}{}'.format(','.join(columns),table,where)

def save_df(df, hdf5_fname, path):

    store = pd.HDFStore(hdf5_fname)

    store[path] = df

    store.close()

    return df

def open_df(hdf5_fname, path):

    store = pd.HDFStore(hdf5_fname)

    df = store[path]

    store.close()

    return df

def is_categorical(var_type):

    return var_type in [variable_type.ORDINAL, variable_type.NOMINAL]

class Bunch(object):

    def __init__(self, **kwds):

        self.__dict__.update(kwds)

def add_subindex(df,subindex_name):

    df = df.sort_index()

    index = df.index

    df[subindex_name] = range(0,len(index))

    duplicated = index.duplicated(keep=False)

    df.loc[duplicated,subindex_name] = df.loc[duplicated].groupby(level=index.names)[subindex_name].apply(lambda x:x - x.min())

    df.loc[~duplicated,subindex_name] = 0

    df.loc[:,subindex_name] = df.loc[:,subindex_name].astype(int)

    df.set_index(subindex_name, append=True,inplace=True)

    return df

def add_same_val_index_level(df,level_val,level_name,axis=0):

    return pd.concat([df], keys=[level_val], names=[level_name],axis=axis)

def set_level_to_same_val(index,level,value):

    return index.set_labels([0]*index.size,level=level).set_levels([value],level=level)

def data_loss(df_start,df_end):

    data_loss = df_start.count().sum() - df_end.count().sum()

    admissions_start = len(df_start.index.get_level_values(column_names.ID).unique().tolist())

    admisstions_end = len(df_end.index.get_level_values(column_names.ID).unique().tolist())

    admission_loss = admissions_start - admisstions_end

    percent_loss = str(round(float(admission_loss)/admissions_start * 100,4))+'% records'

    return df_start.shape,df_end.shape,data_loss,admission_loss,percent_loss

def get_components(df):

    return df.columns.get_level_values('component').unique().tolist()

def filter_columns(df,level,value):

    return df.loc[:,df.columns.get_level_values(level) == value]

def append_to_description(old_desc,addition):

    return old_desc + '(' + addition + ')'

def get_first_dt(df_ts,df_context):

    ids = df_ts.index.get_level_values(column_names.ID).unique()

    admit_dt = get_admit_dt(df_context,ids)

    first_obs_dt = get_first_obs_dt(df_ts)

    first_dt = first_obs_dt.to_frame().join(admit_dt,how='left').apply(pd.np.min,axis=1)

    first_dt.name = column_names.DATETIME

    return first_dt

def get_admit_dt(df_context,ids=ALL):

    if not ids == ALL:

        df_context_filtered = df_context[df_context[column_names.ID].isin(ids)]

    else: df_context_filtered = df_context

    admit_dt = df_context_filtered.loc[:,['id','start_dt']].drop_duplicates(subset=['id']).set_index('id')

    return admit_dt

def get_first_obs_dt(df_ts):

    first_obs_dt = df_ts.groupby(level=column_names.ID).apply(lambda x:x.iloc[0].name[-1])

    first_obs_dt.name = 'start_dt_obs'

    return first_obs_dt

def flatten_index(df,join_char='_',suffix=None,axis=0):

    if axis==0:

        new_vals = []

        for row_name in df.index:

            if type(row_name) is tuple:

                row = join_char.join(map(str,row_name))

            else:

                row = str(row_name)

            if suffix is not None: row = row + join_char + suffix

            new_vals.append(row)

        df.index = pd.Index(new_vals)

    elif axis==1:

        new_vals = []

        for col_name in df.columns:

            if type(col_name) is tuple:

                col = join_char.join(map(str,col_name))

            else:

                col = str(col_name)

            if suffix is not None: col = col + join_char + suffix

            new_vals.append(col)

        df.columns = pd.Index(new_vals)

    else: pass

    return df

def smart_count(col):

    var_type = col.name[2]

    if (var_type == variable_type.NOMINAL) and (col.dtype == pd.np.uint8):

        return col.sum()

    return col.count()

"""

Pytables/HDF5 I/O with axis deconstruction

"""

def read_and_reconstruct(hdf5_fname,path,where=None):

    # Get all paths for dataframes in store

    data_path,col_path = deconstucted_paths(path)

    data = pd.read_hdf(hdf5_fname,data_path,where=where)

    columns = pd.read_hdf(hdf5_fname,col_path)

    return reconstruct_df(data,columns)

def reconstruct_df(data,column_df):

    #reconstruct the columns from dataframe

    column_index = reconstruct_columns(column_df,data.columns)

    df = data.copy()

    df.columns = column_index

    return df

def reconstruct_columns(column_df,col_ix=None):

    column_df = column_df.drop('dtype',axis=1)

    if col_ix is None: col_ix = column_df.index

    col_arys = column_df.loc[col_ix].T.values

    levels = column_df.columns.tolist()

    return pd.MultiIndex.from_arrays(col_arys,names=levels)

def deconstruct_and_write(df,hdf5_fname,path,append=False):

    # Deconstruct the dataframe

    data,columns = deconstruct_df(df)

    # Get all paths for dataframes in store

    data_path,col_path = deconstucted_paths(path)

    #Open store and save df

    store = pd.HDFStore(hdf5_fname)

    store.put(data_path,data,append=append,format='t')

    if (not append) or col_path not in store:

        columns.to_hdf(hdf5_fname,col_path,format='t')

    store.close()

    return

def deconstruct_df(df):

    columns = pd.DataFrame(map(list,df.columns.tolist()),columns=df.columns.names)

    columns['dtype'] = df.dtypes.values.astype(str)

    data = df.copy()

    data.columns = [i for i in range(df.shape[1])]

    return data,columns

def deconstucted_paths(path):

    data_path = '{}/{}'.format(path,'data')

    col_path = '{}/{}'.format(path,'columns')

    return data_path,col_path

def complex_row_mask(df,specs,operator='or'):

    #init our mask datframe with row index of passed in dataframe

    df_mask = pd.DataFrame(index=df.index)

    #make sure our specs are a list

    if not isinstance(specs,list): specs = [specs]

    #if we have no specs, then we are not going to mask anything

    if (specs is None) or (len(specs) == 0):

        df_mask.loc[:,0] = True

    #apply specs to create mask

    for idx,spec in enumerate(specs):

        df_spec_mask = pd.DataFrame(index=df.index)

        for col_name,spec_info in spec.iteritems():

            if callable(spec_info):

                df_spec_mask.loc[:,col_name] = df.loc[:,col_name].apply(spec_info)

            else:

                if not isinstance(spec_info,list): spec_info = [spec_info]

                df_spec_mask.loc[:,col_name] = df.loc[:,col_name].isin(spec_info)

        df_mask.loc[:,idx] = df_spec_mask.all(axis=1)

    #if or, will will include rows from each spec.

    #   if and, only rows that meet criteria of EVERY spec ar included

    if operator == 'or' : mask = df_mask.any(axis=1)

    else: mask = df_mask.all(axis=1)

    return mask

def smart_join(hdf5_fname,paths,joined_path,ids,

                                        chunksize=5000,

                                        need_deconstruct=True,

                                        hdf5_fname_for_join=None,

                                        overwrite=True):

    logger.log('Smart join: n={}, {}'.format(len(ids),paths),new_level=True)

    if hdf5_fname_for_join is None: hdf5_fname_for_join=hdf5_fname

    store = pd.HDFStore(hdf5_fname_for_join)

    if (joined_path in store):

        if overwrite: del store[joined_path]

        else :

            store.close()

            logger.end_log_level()

            return hdf5_fname_for_join

    #sort ids, should speed up where clauses and selects

    ids = sorted(ids)

    #do chunked join

    logger.log('JOINING dataframes',new_level=True)

    for ix_start in range(0,len(ids),chunksize):

        ix_end = min(ix_start + chunksize,len(ids))

        id_slice = ids[ix_start:ix_end]

        where = '{id_col} in {id_list}'.format(id_col=column_names.ID,id_list=id_slice)

        logger.log('Slice & Join: {} --> {}, n={}'.format(id_slice[0], id_slice[-1],len(id_slice)),new_level=True)

        df_slice = None

        # for path in df_dict.keys():

        for path in paths:

            try:

                logger.log(path)

                if need_deconstruct: slice_to_add = read_and_reconstruct(hdf5_fname,path,where=where)

                else: slice_to_add = pd.read_hdf(hdf5_fname,path,where=where)

            except KeyError as err:

                logger.log(end_prev=True,start=False)

                print err

                continue

            if df_slice is None: df_slice = slice_to_add

            else:

                df_slice = df_slice.join(slice_to_add,how='outer')

                del slice_to_add

        logger.end_log_level()

        logger.log('Append slice')

        if need_deconstruct: deconstruct_and_write(df_slice,hdf5_fname_for_join,joined_path,append=True)

        else: df_slice.to_hdf(hdf5_fname_for_join,joined_path,append=True,format='t')

        del df_slice

    logger.end_log_level()

    logger.end_log_level()

    return hdf5_fname_for_join

def make_list_hash(l):

    #need to sort and make sure list are unique before hashing

    l = sorted(list(set(l)))

    #use a hash to make sure store this set of ids uniquely

    key = hash(''.join(map(str,l)))

    return key

"""

Dask intelligent join

"""

def dask_open_and_join(hdf5_fname,path,components,ids=ALL,chunksize=500000):

    df_all=None

    logger.log('DASK OPEN & JOIN n={} components: {}'.format(len(components),components),new_level=True)

    for component in components:

        logger.log('{}: {}/{}'.format(component.upper(),components.index(component)+1,len(components)),new_level=True)

        df_comp = open_df(hdf5_fname,'{}/{}'.format(path,component))

        df_comp.sort_index(inplace=True)

        df_comp.sort_index(inplace=True, axis=1)

        if not ids == ALL:

            df_comp = df_comp[df_comp.index.get_level_values(column_names.ID).isin(ids)]

        logger.log('Convert to dask - {}'.format(df_comp.shape))

        df_dask = dd.from_pandas(df_comp.reset_index(), chunksize=chunksize)

        del df_comp

        logger.log('Join to big DF')

        if df_all is None: df_all = df_dask

        else :

            df_all = df_all.merge(df_dask,how='outer', on=['id','datetime'])

            del df_dask

        logger.end_log_level()

    logger.log('Dask DF back to pandas')

    df_pd = df_all.compute()

    del df_all

    df_pd.set_index(['id','datetime'], inplace=True)

    logger.log('SORT Joined DF')

    df_pd.sort_index(inplace=True)

    df_pd.sort_index(inplace=True, axis=1)

    logger.end_log_level()

    return df_pd

"""

Visualization

"""

import seaborn as sns

import matplotlib.pyplot as plt

def heatmap(df_ts):

    sns.set(context="paper", font="monospace")

    corrmat = df_ts.corr()

    # Set up the matplotlib figure

    f, ax = plt.subplots(figsize=(50, 50))

    # Draw the heatmap using seaborn

    sns.heatmap(corrmat, vmax=1, square=True)

    return