import { join } from 'path'
import { Dataset } from '../../src/dataset/dataset'
import { DatabaseStorage } from '../../src/embeddings/database/storage'
import { BotonicTextClassifier } from '../../src/tasks/text-classification/botonic-text-classifier'
import { TEXT_CLASSIFIER_TEMPLATE } from '../../src/tasks/text-classification/models/types'

const LOCALE = 'en'
const MAX_SEQUENCE_LENGTH = 12

const EMBEDDINGS_DIMENSION = 50
const EMBEDDINGS_TYPE = 'glove'

const EPOCHS = 4
const BATCH_SIZE = 8

const DATASET_DIR_PATH = join('..', 'dataset', LOCALE)

const MODEL_DIR = './model'

const dataset = Dataset.load(DATASET_DIR_PATH)
const { trainSet, testSet } = dataset.split()

;(async () => {
  const classifier = new BotonicTextClassifier(
    LOCALE,
    MAX_SEQUENCE_LENGTH,
    dataset.classes
  )

  classifier.generateVocabulary(trainSet)

  classifier.compile()

  const model = await classifier.createModel(
    TEXT_CLASSIFIER_TEMPLATE.SIMPLE_NN,
    await DatabaseStorage.with(LOCALE, EMBEDDINGS_TYPE, EMBEDDINGS_DIMENSION)
  )

  classifier.importModel(model)

  await classifier.train(trainSet, EPOCHS, BATCH_SIZE)

  const { accuracy, loss } = await classifier.evaluate(testSet)
  console.log(`Test Accuracy: ${accuracy}`)
  console.log(`Test loss: ${loss}`)

  const intents = classifier.classify('I want to return this black shirt')
  console.log(intents)

  await classifier.saveModel(MODEL_DIR)
})()
