import { join } from 'path'

import { Dataset } from '../src/dataset/dataset'
import { DatabaseStorage } from '../src/embeddings/database/storage'
import { BotonicNer } from '../src/tasks/ner/botonic-ner'
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 = 'dataset'
const MODELS_DIR = 'models'
const NER_DIR = 'ner'
const TEXT_CLASSIFICATION_DIR = 'text-classification'

const DATASET_DIR_PATH = join(__dirname, DATASET_DIR, LOCALE)

const NER_DIR_PATH = join(__dirname, MODELS_DIR, NER_DIR)
const TEXT_CLASSIFICATION_DIR_PATH = join(
  __dirname,
  MODELS_DIR,
  TEXT_CLASSIFICATION_DIR
)

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

;(async () => {
  const wordEmbeddingStorage = await DatabaseStorage.with(
    LOCALE,
    EMBEDDINGS_TYPE,
    EMBEDDINGS_DIMENSION
  )

  //NER
  const ner = BotonicNer.with(LOCALE, MAX_SEQUENCE_LENGTH, dataset.entities)
  ner.generateVocabulary(trainSet)
  ner.compile()
  await ner.createModel('biLstm', wordEmbeddingStorage)
  await ner.train(trainSet, EPOCHS, BATCH_SIZE)
  const nerEvaluation = await ner.evaluate(testSet)
  console.log(`Test Accuracy: ${nerEvaluation.accuracy}`)
  console.log(`Test loss: ${nerEvaluation.loss}`)
  await ner.saveModel(NER_DIR_PATH)

  //TEXT CLASSIFIER
  const classifier = new BotonicTextClassifier(
    LOCALE,
    MAX_SEQUENCE_LENGTH,
    dataset.classes
  )
  classifier.generateVocabulary(trainSet)
  classifier.compile()
  await classifier.createModel(
    TEXT_CLASSIFIER_TEMPLATE.SIMPLE_NN,
    wordEmbeddingStorage
  )
  await classifier.train(trainSet, EPOCHS, BATCH_SIZE)
  const classifierEvaluation = await classifier.evaluate(testSet)
  console.log(`Test Accuracy: ${classifierEvaluation.accuracy}`)
  console.log(`Test loss: ${classifierEvaluation.loss}`)
  await classifier.saveModel(TEXT_CLASSIFICATION_DIR_PATH)

  //Testing both models:
  const sentence = 'I want to return this black shirt'
  const intents = classifier.classify(sentence)
  const entities = ner.recognizeEntities(sentence)
  console.log(`Sentence: ${sentence}`)
  console.log(entities)
  console.log(intents)
})()
