import json

import sparse

import pandas as pd

import numpy as np

import scipy.sparse

import joblib

def load_IDs(fname):

    IDs = pd.read_csv(fname, header=0, names=['ID'])

    IDs.index.name = 'i'

    IDs = IDs.reset_index()

    return IDs

def _get_feature_set(df, X_ALL, IDs_ALL):

    IDs = df.set_index('ID')[[]]

    idx = IDs.join(IDs_ALL.set_index('ID')).astype(float)

    X = [X_ALL[int(i),:] if not np.isnan(i) else sparse.zeros(X_ALL.shape[1]) for i in idx.values]

    return sparse.stack(X)

def get_features(df, feature_sets):

    features = []

    feature_names = []

    if 'demog' in feature_sets:

        X_d = df.set_index('hosp_id')[['ID']].join(df_demog).reset_index(drop=True).set_index('ID').loc[df['ID']]

        X_d = sparse.as_coo(X_d.values)

        features.append(X_d)

        feature_names.append(names_demog)

        print('demog - Done')

    if 'vitals' in feature_sets:

        X_v = _get_feature_set(df, X_vitals, IDs_vitals)

        features.append(X_v)

        feature_names.append(names_vitals)

        print('vitals - Done')

    if 'meds' in feature_sets:

        X_m = _get_feature_set(df, X_meds, IDs_meds)

        features.append(X_m)

        feature_names.append(names_meds)

        print('meds - Done')

    if 'labs' in feature_sets:

        X_l = _get_feature_set(df, X_labs, IDs_labs)

        features.append(X_l)

        feature_names.append(names_labs)

        print('labs - Done')

    if 'flow' in feature_sets:

        print('flow', end='')

        X_f = _get_feature_set(df, X_flow, IDs_flow)

        features.append(X_f)

        feature_names.append(names_flow)

        print(' - Done')

    X = sparse.concatenate(features, axis=1).tocsr()

    feature_names = sum(feature_names, [])

    return X, np.array(feature_names)

if __name__ == '__main__':

    df_demog = pd.read_csv('sample_output/out_demog/static-features.csv').set_index('hosp_id')

    names_demog = list(df_demog.columns)

    print('demog - Loaded')

    X_vitals = sparse.load_npz('sample_output/out_vitals/X_all.npz')

    IDs_vitals = load_IDs('sample_output/out_vitals/X_all.IDs.csv')

    names_vitals = json.load(open('metadata/vitals/X_all.feature_names.json', 'r'))

    print('vitals - Loaded')

    X_meds = sparse.load_npz('sample_output/out_meds/X_all.npz')

    IDs_meds = load_IDs('sample_output/out_meds/X_all.IDs.csv')

    names_meds = json.load(open('metadata/meds/X_all.feature_names.json', 'r'))

    print('meds - Loaded')

    X_labs = sparse.load_npz('sample_output/out_labs/X_all.npz')

    IDs_labs = load_IDs('sample_output/out_labs/X_all.IDs.csv')

    names_labs = json.load(open('metadata/labs/X_all.feature_names.json', 'r'))

    print('labs - Loaded')

    X_flow = sparse.load_npz('sample_output/out_flow/X_all.npz')

    IDs_flow = load_IDs('sample_output/out_flow/X_all.IDs.csv')

    names_flow = json.load(open('metadata/flow/X_all.feature_names.json', 'r'))

    print('flow - Loaded')

    df_cohort = pd.read_csv('sample_input/windows_map.csv')

    X, names = get_features(df_cohort, ['demog', 'vitals', 'meds', 'labs', 'flow'])

    df_features = pd.DataFrame(X.todense(), columns=names, index=df_cohort['ID'])

    pd.Series(names).rename('feature_name').to_csv('./sample_output/feature_names.csv', index=False)

    ## Full feature matrix

    joblib.dump(df_features, 'sample_output/full.joblib')

    ## Baseline features

    baseline_cols = pd.read_csv('metadata/Baseline_Feature_Names.txt', sep='\t', header=None)[0].values

    df_baseline = df_features[baseline_cols]

    df_baseline.to_csv('sample_output/baseline.csv')

    ## M-CURES (lite)

    mcures_cols = pd.read_csv('metadata/MCURES_Feature_Names.txt', sep='\t', header=None)[0].values

    df_mcures = df_features[mcures_cols]

    df_mcures.to_csv('sample_output/mcures.csv')