type W2VModel struct {
	Words   int		
	Size    int 			
	Vocab   []string
	WordIdx map[string]int
	Vec     [][]float32
}

func (m *W2VModel) Load(fn string) {
	file, err := os.Open(fn)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Fscanf(file, "%d", &m.Words)
	fmt.Fscanf(file, "%d", &m.Size)

	var ch string
	m.Vocab = make([]string, m.Words)
	m.Vec = make([][]float32, m.Words)
	m.WordIdx = make(map[string]int)
	for b := 0; b < m.Words; b++ {
		m.Vec[b] = make([]float32, m.Size)
		fmt.Fscanf(file, "%s%c", &m.Vocab[b], &ch)
		m.WordIdx[m.Vocab[b]] = b
		binary.Read(file, binary.LittleEndian, m.Vec[b])

		length := 0.0
		for _, v := range m.Vec[b] {
			length += float64(v * v)
		}
		length = math.Sqrt(length)

		for i, _ := range m.Vec[b] {
			m.Vec[b][i] /= float32(length)
		}
	}
	file.Close()
}

type PoemModel struct {
	Poems    []string   `json:"poems"`
	Bags     [][]string `json:"bags"`
	W2V      W2VModel
	Vectors  [][][]float32
	Matrices []mat.Matrix
}

func (pm *PoemModel) LoadJsonModel(fileName string) error {
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}

	err = json.Unmarshal(file, pm)
	if err != nil {
		return err
	}
	return nil
}

func (pm *PoemModel) Matricize() {
	pm.Matrices = make([]mat.Matrix, len(pm.Bags))
	for idx, bag := range pm.Bags {
		data, rows := pm.TokenVectorsData(bag)
		pm.Matrices[idx] = mat.NewDense(rows, pm.W2V.Size, data).T()
	}
}

func (pm *PoemModel) TokenizeWords(words []string) []string {
	POS_TAGS := map[string]string {
		"NOUN": "_NOUN",
		"VERB": "_VERB", "INFN": "_VERB", "GRND": "_VERB", "PRTF": "_VERB", "PRTS": "_VERB",
		"ADJF": "_ADJ", "ADJS": "_ADJ",
		"ADVB": "_ADV",
		"PRED": "_ADP",
	}

	STOP_TAGS := map[string]bool {"PREP": true, "CONJ": true, "PRCL": true, "NPRO": true, "NUMR": true}

	result := make([]string, 0, len(words))

	for _, w := range words {
		_, morphNorms, morphTags := morph.Parse(w)
		if len(morphNorms) == 0 {
			continue
		}

		suffixes := make(map[string]bool) // added suffixes

		for i, tags := range morphTags {
			norm := morphNorms[i]
			tag := strings.Split(tags, ",")[0]
			_, hasStopTag := STOP_TAGS[tag]
			if hasStopTag {
				break
			}

			suffix, hasPosTag := POS_TAGS[tag]
			_, hasSuffix := suffixes[suffix]
			if hasPosTag && ! hasSuffix {
				result = append(result, norm + suffix)
				suffixes[suffix] = true
			}
		}
	}

	return result
}

func (pm *PoemModel) SimilarPoemsMx(queryWords []string, topN int) []string {
	simPoems := make([]string, 0, topN)
	tokens := pm.TokenizeWords(queryWords)
	queryData, queryVecsN := pm.TokenVectorsData(tokens)
	if len(tokens) == 0 || topN <= 0 || queryVecsN == 0{
		return simPoems
	}

	queryMx := mat.NewDense(queryVecsN, pm.W2V.Size, queryData)

	type PoemSimilarity struct {
		Idx	int
		Sim float64
	}

	sims := make([]PoemSimilarity, len(pm.Bags))

	for idx, _ := range pm.Bags {
		var resMx mat.Dense
		bagMx := pm.Matrices[idx]
		_, poemVecsN := bagMx.Dims()
		resMx.Mul(queryMx, bagMx)
		sim := mat.Sum(&resMx)

		if poemVecsN > 0 {
			sim /= float64(poemVecsN + queryVecsN)
		}

		sims[idx].Idx = idx
		sims[idx].Sim = sim
	}

	sort.Slice(sims, func (i, j int) bool {
		return sims[i].Sim > sims[j].Sim
	})

	for i := 0; i < topN; i ++ {
		simPoems = append(simPoems, pm.Poems[sims[i].Idx])
	}

	return simPoems
}