import json

import sparse

import pandas as pd

import numpy as np

import joblib

if __name__ == '__main__':

    df_static = pd.read_csv('sample_input/demog.csv').set_index('hosp_id')

    discretization_bins = json.load(open('metadata/demog/discretization.json', 'r'))

    feature_names = json.load(open('metadata/demog/s_all.feature_names.json', 'r'))

    cols_useful = list(discretization_bins.keys())

    df_static = df_static[cols_useful]

    ##### FIDDLE

    from FIDDLE_steps_2 import *

    from FIDDLE_helpers import *

    df = df_static

    df_by_cols = [df[col] for col in df.columns]

    out_0 = [smart_qcut_bins((x, discretization_bins[x.name])) for x in tqdm(df_by_cols, desc='pd.qcut')]

    cols_data, dis_bins = zip(*out_0)

    out = [dummify(z) for z in tqdm(cols_data, desc='dummify')]

    df_features = pd.concat(out, axis=1).sort_index(axis=1)

    # Drop any values='missing'

    df_features = df_features.loc[:, [col for col in df_features.columns if 'missing' not in col]]

    # Make sure same features

    assert set(df_features.columns) <= set(feature_names)

    df_features = df_features.reindex(columns=feature_names, fill_value=0)

    sdf = df_features.astype(pd.SparseDtype(int, fill_value=0))

    X_all = sparse.COO(sdf.sparse.to_coo())

    sparse.save_npz('sample_output/out_demog/X_all.npz', X_all)

    df_features[[]].to_csv('sample_output/out_demog/X_all.IDs.csv')

    import sparse, json

    s = sparse.load_npz('sample_output/out_demog/X_all.npz').todense()

    cols = feature_names

    df_s = pd.DataFrame(s, columns=cols)

    df_s.index = df_static.index

    df_s.to_csv('sample_output/out_demog/static-features.csv')