import os, sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import reader, submitted, json, argparse
import numpy as np

#################################################################
# visible params
TRAIN_DIR        ='data/train'                  # directory containing train data
DEV_DIR          ='data/dev'                    # directory containing dev data
MULTICLASS_DATA  ='human'                       # which parquet file to run multiclass test on
OUTFILE          ='solution.json'               # filename in which to store results

# # hidden params
# TRAIN_DIR        ='data/hidden/train'           # directory containing train data
# DEV_DIR          ='data/hidden/dev'             # directory containing dev data
# MULTICLASS_DATA  ='llm'                         # which parquet file to run multiclass test on
# OUTFILE          ='solution_hidden.json'               # filename in which to store results

# common params
SMOOTHNESS       =1                             # preset Laplace smoothing parameter
PRIOR            =[1.0/6]*6                     # preset prior
SMOOTHNESS_SPEC  =[0.01, 0.001]                 # list of smoothnesses
PRIOR_SPEC       =[[1.0/6]*6, [0.1]*5+[0.5]]    # list of priors

# Run submitted code (multiclass)
train = {
    'acad': reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['acad']),
    'blog': reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['blog']),
    'news': reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['news']),
    'fic' : reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['fic']),
    'spok': reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['spok']),
    'tvm' : reader.loadDataset(f'{TRAIN_DIR}/{MULTICLASS_DATA}.parquet', ['tvm']),
}
val = {
    'acad': reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['acad']),
    'blog': reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['blog']),
    'news': reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['news']),
    'fic' : reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['fic']),
    'spok': reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['spok']),
    'tvm' : reader.loadDataset(f'{DEV_DIR}/{MULTICLASS_DATA}.parquet', ['tvm']),
}
labels, texts = reader.val_set_helper(val)

CLASSES = ['acad', 'blog', 'news']
frequency = submitted.create_frequency_table(train)
nonstop = submitted.remove_stopwords(frequency)
likelihood = submitted.laplace_smoothing(nonstop, SMOOTHNESS)
hypotheses = submitted.naive_bayes(texts, likelihood, PRIOR, CLASSES)
accuracies = submitted.optimize_hyperparameters(texts, labels, CLASSES, nonstop, PRIOR_SPEC, SMOOTHNESS_SPEC)

# Package it all into a dictionary, and save to a JSON file
solution = {
    'mc_prior' : PRIOR,
    'mc_smoothness' : SMOOTHNESS,
    'mc_priors' : PRIOR_SPEC,
    'mc_smoothnesses' : SMOOTHNESS_SPEC,
    'mc_frequency' : frequency,
    'mc_nonstop' : nonstop,
    'mc_likelihood' : likelihood,
    'mc_hypotheses' : hypotheses,
    'mc_accuracies' : [ [ a for a in accuracies[m] ] for m in range(accuracies.shape[0]) ],
    'mc_classes' : CLASSES
}
with open(OUTFILE, 'w') as f:
    json.dump(solution, f, indent=4)
