# coding: utf-8

import os, io
import argparse
import funcy
import json
import numpy as np
from pathlib import Path

from tmnt.inference import BowVAEInferencer

parser = argparse.ArgumentParser()

parser.add_argument('--model_dir', type=Path, help='Directory with trained model files')
parser.add_argument('--vec_file', type=str, help='Sparse vector file')
parser.add_argument('--num_terms', type=int, help='Number of terms to display for each topic (for each covariate value)')
parser.add_argument('--label_map', type=str, help='Label map generated by vectorizer', default=None)
parser.add_argument('--covariate_values', type=str, help='File with covariate values in ordered list')
parser.add_argument('--scalar_covar_range', type=str, help='min/max comma-separated (default [0.0, 1.0])', default="0.0,1.0")
parser.add_argument('--num_covars', type=int, help='Number of points within interval', default=0)
parser.add_argument('--output_file', type=str, help='Text file output')



def print_categorical_covar_model(inferencer, label_map, n_terms, output_file):    
    topk_per_covar = inferencer.get_top_k_words_per_topic_per_covariate(n_terms)
    inv_map = sorted( [(v,k) for k,v in label_map.items()] )
    with open(output_file, 'w') as out:
        for (ind,n) in inv_map:
            out.write("===== Topics for {} =====\n".format(n))
            c_topics = topk_per_covar[ind]
            i = 0
            for topic in c_topics:
                i += 1
                out.write("Topic {}: ".format(str(i)))
                for term in topic:
                    out.write("{} ".format(term))
                out.write('\n')

def read_label_map(label_map):
    with io.open(label_map, 'r') as fp:
        lm_dict = json.load(fp)
    return lm_dict

if __name__ == '__main__':
    args = parser.parse_args()
    inferencer = BowVAEInferencer.from_saved(model_dir=args.model_dir)
    label_map = read_label_map(args.label_map)
    print_categorical_covar_model(inferencer, label_map, args.num_terms, args.output_file)
