package com.davezhao.odds

import com.davezhao.models.Connector
import com.davezhao.models.StdQaInfoTable
import com.davezhao.nlp.WordBreaker
import com.hankcs.hanlp.HanLP
import com.hankcs.hanlp.mining.word2vec.Word2VecTrainer
import org.jetbrains.exposed.sql.selectAll
import org.jetbrains.exposed.sql.transactions.transaction
import org.slf4j.LoggerFactory
import java.io.BufferedWriter
import java.io.File
import java.io.FileWriter

object TrainWoedVec {
    val log = LoggerFactory.getLogger(this::class.java)

    fun segmentWord(rawFilepath: String): String {
        val outFilepath = "./seg_data.csv"
        val fout = BufferedWriter(FileWriter(File(outFilepath)))
        File(rawFilepath).readLines().forEach { line ->
            val trim = line.trim()
            if (trim.isNotEmpty() && !trim.startsWith("#")) {
                val segLine = WordBreaker.segment(line).map { seg ->
                    seg.word
                }/*.filter { word -> !StopWords.data.contains(word) }*/ .joinToString(" ")
                fout.write(segLine + "\n")
            }
        }

        fout.close()

        return outFilepath
    }

    fun trainModel(dataPath: String) {
        val trainerBuilder = Word2VecTrainer();
        trainerBuilder.setNumIterations(50);
        trainerBuilder.setLayerSize(100);
        trainerBuilder.useNumThreads(4)
        val wordVectorModel = trainerBuilder.train(dataPath, "./word_vector_model.bin")

    }

    fun pullTrainedData() {
        Connector.open()
        val fout = BufferedWriter(FileWriter("./data/w2v/seg_data.txt"))
        transaction {
            StdQaInfoTable.selectAll().forEach {
                val coreMsg = it[StdQaInfoTable.question] + " " + it[StdQaInfoTable.answer] + " " + it[StdQaInfoTable.evidence]
                val seggedSentence = HanLP.segment(coreMsg).map { it.word } .joinToString(" ")
                log.info(seggedSentence)
                fout.write(seggedSentence + "\n")
            }
        }
        fout.close()
    }
}

fun main(args: Array<String>) {
//    TrainWoedVec.segmentWord("./raw_data_0125.csv")
//    TrainWoedVec.trainModel("./seg_data.csv")
    TrainWoedVec.pullTrainedData()
}