import numpy as np
import pandas as pd
import savReaderWriter as sr
import pyreadstat

from collections import defaultdict
from quantipy.core.tools.dp.prep import start_meta, condense_dichotomous_set
import os

def parse_sav_file(filename, path=None, name="", ioLocale="en_US.UTF-8", ioUtf8=True, dichot=None,
                   dates_as_strings=False, text_key="en-GB", engine='savReaderWriter'):
    """ Parses a .sav file and returns a touple of Data and Meta

        Parameters
        ----------
        filename : str, name of sav file
            path : str, the path to the sav file
            name : str, a name for the sav (stored in the meta)
        ioLocale : str, the locale that SavReaderWriter uses
          ioUtf8 : bool, Boolean that indicates the mode in which text
                         communicated to or from the I/O Module will be.
          dichot : dicit, default=None
                   The values to use for True/False in dichotomous sets
dates_as_strings : bool, default=False
                   If True then all dates from the input SAV will be treated as
                   Quantipy strings.
        text_key : str, default="main"
                   The text_key that all labels should be stored under.

        Returns
        -------
        (data, meta) : The data is a Pandas Dataframe
                     : The meta is a JSON dictionary
    """
    filepath="{}{}".format(path or '', filename)
    filepath = os.path.abspath(filepath)
    data = extract_sav_data(filepath, ioLocale=ioLocale, ioUtf8=ioUtf8, engine=engine)
    meta, data = extract_sav_meta(filepath, name="", data=data, ioLocale=ioLocale,
                                  ioUtf8=ioUtf8, dichot=dichot, dates_as_strings=dates_as_strings,
                                  text_key=text_key, engine=engine)
    return (meta, data)

def extract_sav_data(sav_file, ioLocale='en_US.UTF-8', ioUtf8=True, engine='savReaderWriter'):

    """ see parse_sav_file doc """
    if engine == 'savReaderWriter':
        with sr.SavReader(sav_file, returnHeader=True, ioLocale=ioLocale, ioUtf8=ioUtf8) as reader:
            thedata = [x for x in reader]
            header = thedata[0]
            dataframe = pd.DataFrame.from_records(thedata[1:], coerce_float=False)
            dataframe.columns = header
            for column in header:
                if isinstance(dataframe[column].dtype, np.object):
                    # Replace None with NaN because SRW returns None if casting dates fails (dates are of type np.object))
                    values = dataframe[column].dropna().values
                    if len(values) > 0:
                        if isinstance(values[0], str):
                            dataframe[column] = dataframe[column].dropna().map(str.strip)
                        elif isinstance(values[0], str):
                            # savReaderWriter casts dates to str
                            dataframe[column] = dataframe[column].dropna().map(str.strip)
                            # creating DATETIME objects should happen here
            return dataframe
    elif engine == 'readstat':
        df, meta = pyreadstat.read_sav(sav_file)
        return df

def extract_sav_meta(sav_file, name="", data=None, ioLocale='en_US.UTF-8',
                     ioUtf8=True, dichot=None, dates_as_strings=False,
                     text_key="en-GB", engine='savReaderWriter'):

    if engine == 'readstat':
        df, metadata = pyreadstat.read_sav(sav_file, encoding=ioLocale.split(".")[-1], metadataonly=True)
        meta = start_meta(text_key=text_key)

        meta['info']['text'] = 'Converted from SAV file {}.'.format(name)
        meta['info']['from_source'] = {'pandas_reader':'sav'}
        meta['sets']['data file']['items'] = [
            'columns@{}'.format(varName)
            for varName in metadata.column_names]

        for index, column in enumerate(metadata.column_names):
            meta['columns'][column] = {}
            meta['columns'][column]['name'] = column
            meta['columns'][column]['parent'] = {}
            if column in metadata.variable_value_labels:
                meta['columns'][column]['values'] = []
                meta['columns'][column]['type'] = "single"
                for value, text in metadata.variable_value_labels[column].items():
                    values = {'text': {text_key: str(text)},
                            'value': int(value)}
                    meta['columns'][column]['values'].append(values)
                    # if user has stored single answer data as a string rather than number
                    # we convert it to floats and store non convertables as nan (with coerce)
                    if column in data.columns and data[column].dtype == 'O':
                        data[column] = pd.to_numeric(data[column], errors='coerce', downcast='float')
            else:
                if column in metadata.original_variable_types:
                    f = metadata.original_variable_types[column]
                    if 'DATETIME' in f:
                        if dates_as_strings:
                            # DATETIME fields from SPSS are currently
                            # being read in as strings because there's an
                            # as-yet undetermined discrepancy between the
                            # input and output dates if datetime64 is used
                            meta['columns'][column]['type'] = 'string'
                        else:
                            meta['columns'][column]['type'] = 'date'
                            data[column] = pd.to_datetime(data[column])
                    elif f.startswith('A'):
                        meta['columns'][column]['type'] = 'string'
                    elif '.' in f:
                        meta['columns'][column]['type'] = "float"
                    else:
                        meta['columns'][column]['type'] = "int"

            # add the variable label to the meta
            meta['columns'][column]['text'] = {text_key : metadata.column_labels[index]}    
        return meta, data

    elif engine == 'savReaderWriter':
        if dichot is None: dichot = {'yes': 1, 'no': 0}

        """ see parse_sav_file doc """
        with sr.SavHeaderReader(sav_file, ioLocale=ioLocale, ioUtf8=ioUtf8) as header:
            # Metadata Attributes
            # ['valueLabels', 'varTypes', 'varSets', 'varAttributes', 'varRoles',
            #  'measureLevels', 'caseWeightVar', 'varNames', 'varLabels', 'formats',
            #  'multRespDefs', 'columnWidths', 'fileAttributes', 'alignments',
            #  'fileLabel', 'missingValues']
            metadata = header.dataDictionary(True)

        meta = start_meta(text_key=text_key)
        meta['info']['text'] = 'Converted from SAV file {}.'.format(name)
        meta['info']['from_source'] = {'pandas_reader':'sav'}
        meta['sets']['data file']['items'] = [
            'columns@{}'.format(varName)
            for varName in metadata.varNames]

        # This should probably be somewhere in the metadata
        # weight_variable_name = metadata.caseWeightVar

        # Descriptions of attributes in metadata are are located here :
        # http://pythonhosted.org/savReaderWriter/#savwriter-write-spss-system-files
        for column in metadata.varNames:
            meta['columns'][column] = {}
            meta['columns'][column]['name'] = column
            meta['columns'][column]['parent'] = {}
            if column in metadata.valueLabels:
                # ValueLabels is type = 'single' (possibry 1-1 map)
                meta['columns'][column]['values'] = []
                meta['columns'][column]['type'] = "single"
                for value, text in metadata.valueLabels[column].items():
                    values = {'text': {text_key: str(text)},
                            'value': int(value)}
                    meta['columns'][column]['values'].append(values)
            else:
                if column in metadata.formats:
                    f = metadata.formats[column]
                    if 'DATETIME' in f:
                        if dates_as_strings:
                            # DATETIME fields from SPSS are currently
                            # being read in as strings because there's an
                            # as-yet undetermined discrepancy between the
                            # input and output dates if datetime64 is used
                            meta['columns'][column]['type'] = 'string'
                        else:
                            meta['columns'][column]['type'] = 'date'
                            data[column] = pd.to_datetime(data[column])
                    elif f.startswith('A'):
                        meta['columns'][column]['type'] = 'string'
                    elif '.' in f:
                        meta['columns'][column]['type'] = "float"
                    else:
                        meta['columns'][column]['type'] = "int"
                else:
                    # Infer meta from data
                    if data is not None:
                        # print "VAR '{}' NOT IN value_labels".format(column)
                        column_values = data[column].dropna()
                        if len(column_values) > 0:
                            # Get the first "not nan" value from the column
                            value = column_values.values[0]
                            if isinstance(value, pd.np.float64):
                                # Float AND Int because savReaderWriter loads them both as float64
                                meta['columns'][column]['text'] = {text_key: [column]}
                                meta['columns'][column]['type'] = "float"
                                if (data[column].dropna() % 1).sum() == 0:
                                    if (data[column].dropna() % 1).unique() == [0]:
                                        try:
                                            data[column] = data[column].astype('int')
                                        except:
                                            pass
                                        meta['columns'][column]['type'] = "int"

                            elif isinstance(value, str) or isinstance(value, str):
                                # Strings
                                meta['columns'][column]['text'] = {text_key: [column]}
                                meta['columns'][column]['type'] = "string"

            if column in metadata.varTypes:
                pass

            if column in metadata.varSets:
                pass

            if column in metadata.varAttributes:
                pass

            if column in metadata.varRoles:
                pass

            if column in metadata.measureLevels:
                pass

            # Some labels are empty strings.note
            if column in metadata.varLabels:
                meta['columns'][column]['text'] = {text_key: metadata.varLabels[column]}

        for mrset in metadata.multRespDefs:
            # meta['masks'][mrset] = {}
            # 'D' is "multiple dichotomy sets" in SPSS
            # 'C' is "multiple category sets" in SPSS
            varNames = list(metadata.multRespDefs[mrset]['varNames'])
            # Find the index where there delimited set should be inserted
            # into data, which is immediately prior to the start of the
            # dichotomous set columns
            dls_idx = data.columns.tolist().index(varNames[0])
            if metadata.multRespDefs[mrset]['setType'] == 'C':
                # Raise if value object of columns is not equal
                if not all(meta['columns'][v]['values'] == meta['columns'][varNames[0]]['values']
                        for v in varNames):
                    msg = 'Columns must have equal values to be combined in a set: {}'
                    raise ValueError(msg.format(varNames))
                # Concatenate columns to set
                df_str = data[varNames].astype('str')
                dls = df_str.apply(lambda x: ';'.join([
                    v.replace('.0', '') for v in x.tolist()
                    if not v in ['nan', 'None']]),
                    axis=1) + ';'
                dls.replace({';': np.NaN}, inplace=True)
                # Get value object
                values = meta['columns'][varNames[0]]['values']

            elif metadata.multRespDefs[mrset]['setType'] == 'D':
                # Generate the delimited set from the dichotomous set
                dls = condense_dichotomous_set(data[varNames], values_from_labels=False, **dichot)
                # Get value object
                values = [{
                            'text': {text_key: metadata.varLabels[varName]},
                            'value': int(v)
                        }
                        for v, varName in enumerate(varNames, start=1)]
            else:
                continue
            # Insert the delimited set into data
            data.insert(dls_idx, mrset, dls)
            # Generate the column meta for the new delimited set
            meta['columns'][mrset] = {
                'name': mrset,
                'type': 'delimited set',
                'text': {text_key: metadata.multRespDefs[mrset]['label']},
                'parent': {},
                'values': values}
            # Add the new delimited set to the 'data file' set
            df_items = meta['sets']['data file']['items']
            df_items.insert(
                df_items.index('columns@{}'.format(varNames[0])),
                'columns@{}'.format(mrset))

            data = data.drop(varNames, axis=1)
            for varName in varNames:
                df_items.remove('columns@{}'.format(varName))
                del meta['columns'][varName]

        return meta, data
