package main

import (
	"bufio"
	"encoding/csv"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// ParsedData представляет структуру для входных и выходных данных
type ParsedData struct {
	Input  string
	Output string
}

// Константы для ограничения количества сообщений и размера словаря
const (
	maxTrainMessages = 50000 // Максимальное количество обучающих сообщений
	maxVocabSize     = 10000 // Максимальный размер словаря для улучшения производительности
)

// TFIDFData хранит предвычисленные TF-IDF векторы и ответы
type TFIDFData struct {
	Vectors [][]float32
	Outputs []string
}

// TFIDF представляет модель TF-IDF
type TFIDF struct {
	Vocab         map[string]int
	IDF           []float64
	DocumentCount int
}

// NewTFIDF инициализирует модель TF-IDF
func NewTFIDF(vocab map[string]int) *TFIDF {
	return &TFIDF{
		Vocab:         vocab,
		IDF:           make([]float64, len(vocab)),
		DocumentCount: 0,
	}
}

// Fit вычисляет IDF по данным
func (tf *TFIDF) Fit(data []ParsedData) {
	df := make([]int, len(tf.Vocab))
	tf.DocumentCount = len(data)

	for _, item := range data {
		tokens := Tokenize(item.Input)
		uniqueTokens := make(map[string]struct{})
		for _, token := range tokens {
			if _, exists := tf.Vocab[token]; exists {
				uniqueTokens[token] = struct{}{}
			}
		}
		for token := range uniqueTokens {
			idx := tf.Vocab[token]
			df[idx]++
		}
	}

	for i := 0; i < len(tf.IDF); i++ {
		tf.IDF[i] = math.Log(float64(tf.DocumentCount) / (1.0 + float64(df[i])))
	}
}

// Transform преобразует текст в TF-IDF вектор
func (tf *TFIDF) Transform(text string) []float32 {
	tokens := Tokenize(text)
	tfidf := make([]float32, len(tf.Vocab))
	tfCounts := make(map[int]float64)

	for _, token := range tokens {
		if idx, exists := tf.Vocab[token]; exists {
			tfCounts[idx] += 1.0
		}
	}

	for idx, count := range tfCounts {
		tfidf[idx] = float32(count * tf.IDF[idx])
	}

	// Нормализация вектора
	var norm float64
	for _, val := range tfidf {
		norm += float64(val * val)
	}
	norm = math.Sqrt(norm)
	if norm == 0 {
		return tfidf
	}
	for i := range tfidf {
		tfidf[i] = float32(float64(tfidf[i]) / norm)
	}

	return tfidf
}

// LoadFiles загружает файлы из указанной директории
func LoadFiles(dirname string, maxMessages int) ([]ParsedData, error) {
	var data []ParsedData
	err := filepath.WalkDir(dirname, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() || len(data) >= maxMessages {
			return nil
		}

		if strings.HasSuffix(d.Name(), ".csv") {
			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer f.Close()

			r := csv.NewReader(f)
			records, err := r.ReadAll()
			if err != nil {
				return err
			}

			for i := 1; i < len(records) && len(data) < maxMessages; i++ { // Предполагается, что первая строка - заголовок
				record := records[i]
				if len(record) >= 2 {
					data = append(data, ParsedData{
						Input:  record[0],
						Output: record[1],
					})
				}
			}
		} else if strings.HasSuffix(d.Name(), ".txt") {
			content, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}
			lines := strings.Split(string(content), "\n")
			for _, line := range lines {
				if len(data) >= maxMessages {
					break
				}
				parts := strings.Split(line, ":")
				if len(parts) > 1 {
					data = append(data, ParsedData{
						Input:  strings.TrimSpace(parts[0]),
						Output: strings.TrimSpace(parts[1]),
					})
				}
			}
		}
		return nil
	})
	return data, err
}

// Tokenize разбивает текст на токены
func Tokenize(text string) []string {
	return strings.Fields(strings.ToLower(strings.TrimSpace(text))) // Приводим к нижнему регистру для консистентности
}

// BuildVocab строит словарь из набора данных
func BuildVocab(data []ParsedData) map[string]int {
	wordFreq := make(map[string]int)

	for _, item := range data {
		tokens := append(Tokenize(item.Input), Tokenize(item.Output)...)
		for _, token := range tokens {
			wordFreq[token]++
		}
	}

	type kv struct {
		Key   string
		Value int
	}

	var freqs []kv
	for k, v := range wordFreq {
		freqs = append(freqs, kv{k, v})
	}

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

	if len(freqs) > maxVocabSize {
		freqs = freqs[:maxVocabSize]
	}

	vocab := make(map[string]int)
	for i, kv := range freqs {
		vocab[kv.Key] = i
	}
	return vocab
}

// SaveVocabulary сохраняет словарь в файл
func SaveVocabulary(vocab map[string]int, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	encoder := gob.NewEncoder(f)
	return encoder.Encode(vocab)
}

// LoadVocabulary загружает словарь из файла
func LoadVocabulary(filename string) (map[string]int, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var vocab map[string]int
	decoder := gob.NewDecoder(f)
	if err := decoder.Decode(&vocab); err != nil {
		return nil, err
	}
	return vocab, nil
}

// SaveTFIDF сохраняет модель TF-IDF в файл
func SaveTFIDF(tfidf *TFIDF, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	encoder := gob.NewEncoder(f)
	return encoder.Encode(tfidf)
}

// LoadTFIDF загружает модель TF-IDF из файла
func LoadTFIDF(filename string) (*TFIDF, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var tfidf TFIDF
	decoder := gob.NewDecoder(f)
	if err := decoder.Decode(&tfidf); err != nil {
		return nil, err
	}
	return &tfidf, nil
}

// SaveTFIDFData сохраняет предвычисленные TF-IDF векторы и ответы в файл
func SaveTFIDFData(data *TFIDFData, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	encoder := gob.NewEncoder(f)
	return encoder.Encode(data)
}

// LoadTFIDFData загружает предвычисленные TF-IDF векторы и ответы из файла
func LoadTFIDFData(filename string) (*TFIDFData, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var data TFIDFData
	decoder := gob.NewDecoder(f)
	if err := decoder.Decode(&data); err != nil {
		return nil, err
	}
	return &data, nil
}

// CosineSimilarity вычисляет косинусное сходство между двумя векторами
func CosineSimilarity(a, b []float32) float64 {
	if len(a) != len(b) {
		log.Fatalf("Vectors have different lengths: %d vs %d", len(a), len(b))
	}
	var dotProduct float64
	var normA float64
	var normB float64
	for i := 0; i < len(a); i++ {
		dotProduct += float64(a[i] * b[i])
		normA += float64(a[i] * a[i])
		normB += float64(b[i] * b[i])
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
}

// trainModel подготавливает и сохраняет TF-IDF векторы и ответы
func trainModel(trainDir string, maxTrainMessages int) {
	parsedData, err := LoadFiles(trainDir, maxTrainMessages)
	if err != nil {
		log.Fatalf("Ошибка при загрузке обучающих файлов: %v", err)
	}

	fmt.Printf("Загружено %d обучающих записей\n", len(parsedData))

	vocab := BuildVocab(parsedData)

	// Создание директории models, если не существует
	if err := os.MkdirAll("models", os.ModePerm); err != nil {
		log.Fatalf("Ошибка при создании директории models: %v", err)
	}

	// Сохранение словаря
	if err := SaveVocabulary(vocab, "models/vocab.gob"); err != nil {
		log.Fatalf("Ошибка при сохранении словаря: %v", err)
	}

	// Создание и обучение TF-IDF модели
	tfidf := NewTFIDF(vocab)
	tfidf.Fit(parsedData)
	if err := SaveTFIDF(tfidf, "models/tfidf.gob"); err != nil {
		log.Fatalf("Ошибка при сохранении модели TF-IDF: %v", err)
	}

	// Преобразование всех входов в TF-IDF векторы
	fmt.Println("Преобразование обучающих данных в TF-IDF векторы...")
	trainVectors := make([][]float32, len(parsedData))
	outputs := make([]string, len(parsedData))
	for i, item := range parsedData {
		trainVectors[i] = tfidf.Transform(item.Input)
		outputs[i] = item.Output
		if (i+1)%10000 == 0 {
			fmt.Printf("Преобразовано %d/%d записей\n", i+1, len(parsedData))
		}
	}

	// Сохранение предвычисленных TF-IDF векторов и соответствующих ответов
	data := TFIDFData{
		Vectors: trainVectors,
		Outputs: outputs,
	}
	if err := SaveTFIDFData(&data, "models/tfidf_data.gob"); err != nil {
		log.Fatalf("Ошибка при сохранении TF-IDF данных: %v", err)
	}

	fmt.Println("TF-IDF векторы и ответы успешно сохранены.")
}

// interactWithModel позволяет взаимодействовать с обученной моделью
func interactWithModel() {
	// Загрузка модели TF-IDF
	tfidf, err := LoadTFIDF("models/tfidf.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке модели TF-IDF: %v", err)
	}

	// Загрузка предвычисленных TF-IDF векторов и ответов
	tfidfData, err := LoadTFIDFData("models/tfidf_data.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке TF-IDF данных: %v", err)
	}

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("Задайте вопрос (или введите 'exit' для выхода): ")
		text, err := reader.ReadString('\n')
		if err != nil {
			log.Fatalf("Ошибка при чтении ввода: %v", err)
		}
		text = strings.TrimSpace(text)

		if strings.ToLower(text) == "exit" {
			break
		}

		// Преобразование пользовательского ввода в TF-IDF вектор
		inputVec := tfidf.Transform(text)

		// Вычисление косинусного сходства с каждым обучающим вектором
		maxSim := -1.0
		bestIdx := -1
		for i, trainVec := range tfidfData.Vectors {
			sim := CosineSimilarity(inputVec, trainVec)
			if sim > maxSim {
				maxSim = sim
				bestIdx = i
			}
		}

		var response string
		if bestIdx != -1 && maxSim > 0.1 { // Порог сходства можно настроить
			response = tfidfData.Outputs[bestIdx]
		} else {
			response = "Извините, я не понял ваш вопрос."
		}

		fmt.Printf("Ответ: %s (Сходство: %.4f)\n", response, maxSim)
	}
}

// validateModel выполняет валидацию модели на данных из validationDir
func validateModel(validationDir string) {
	// Загрузка модели TF-IDF
	tfidf, err := LoadTFIDF("models/tfidf.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке модели TF-IDF: %v", err)
	}

	// Загрузка предвычисленных TF-IDF векторов и ответов
	tfidfData, err := LoadTFIDFData("models/tfidf_data.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке TF-IDF данных: %v", err)
	}

	// Загрузка валидационных данных
	validationData, err := LoadFiles(validationDir, maxTrainMessages)
	if err != nil {
		log.Fatalf("Ошибка при загрузке валидационных файлов: %v", err)
	}

	fmt.Printf("Загружено %d валидационных записей\n", len(validationData))

	correct := 0
	total := len(validationData)

	for idx, item := range validationData {
		// Преобразование входного текста в TF-IDF вектор
		inputVec := tfidf.Transform(item.Input)

		// Вычисление косинусного сходства с каждым обучающим вектором
		maxSim := -1.0
		bestIdx := -1
		for i, trainVec := range tfidfData.Vectors {
			sim := CosineSimilarity(inputVec, trainVec)
			if sim > maxSim {
				maxSim = sim
				bestIdx = i
			}
		}

		var predictedOutput string
		if bestIdx != -1 && maxSim > 0.1 { // Порог сходства можно настроить
			predictedOutput = tfidfData.Outputs[bestIdx]
		} else {
			predictedOutput = "Извините, я не понял ваш вопрос."
		}

		// Проверка корректности
		if strings.ToLower(predictedOutput) == strings.ToLower(item.Output) {
			correct++
		}

		// Вывод прогресса
		if (idx+1)%1000 == 0 || idx+1 == total {
			fmt.Printf("Валидировано %d/%d записей\n", idx+1, total)
		}
	}

	accuracy := float64(correct) / float64(total) * 100
	fmt.Printf("Validation Accuracy: %.2f%%\n", accuracy)
}

// main - точка входа программы
func main() {
	fmt.Print("Enter mode of operation (t for training, v for validation, d for dialogue): ")
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	mode := scanner.Text()

	switch mode {
	case "t":
		trainModel("train", maxTrainMessages)
	case "v":
		validateModel("valid")
	case "d":
		interactWithModel()
	default:
		fmt.Println("Unknown mode. Use 't' for training, 'v' for validation, or 'd' for dialogue.")
	}
}


Enter mode of operation (t for training, v for validation, d for dialogue): v
Загружено 3727 валидационных записей

очень долго обрабатывает исправь