package main

import (
	"bufio"
	"encoding/csv"
	"encoding/gob"
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"

	"gorgonia.org/gorgonia"
	"gorgonia.org/tensor"
)

type ParsedData struct {
	Input  string
	Output string
}

const (
	maxTrainMessages = 1000000 // Максимальное количество обучающих сообщений
	maxVocabSize     = 10000   // Уменьшено до 10,000 для эффективности
	numEpochs        = 10      // Количество эпох обучения
	learningRate     = 0.001   // Скорость обучения
	batchSize        = 150     // Размер батча
	embeddingSize    = 100     // Размер эмбеддингов
	maxSeqLength     = 10      // Максимальная длина последовательности
)

// LoadCSVFiles загружает и парсит CSV файлы из указанной директории
func LoadCSVFiles(dirname string) ([]ParsedData, error) {
	files, err := os.ReadDir(dirname)
	if err != nil {
		return nil, err
	}

	var data []ParsedData
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".csv") {
			f, err := os.Open(filepath.Join(dirname, file.Name()))
			if err != nil {
				return nil, err
			}
			defer f.Close()

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

			for i := 1; i < len(records); i++ { // Предполагается, что первая строка — заголовок
				record := records[i]
				if len(record) >= 6 {
					data = append(data, ParsedData{
						Input:  record[4],
						Output: record[5],
					})
				}
			}
		}
	}
	return data, nil
}

// Tokenize разбивает текст на токены
func Tokenize(text string) []string {
	reg, err := regexp.Compile("[^a-zA-Z0-9]+")
	if err != nil {
		log.Fatalf("Regex compilation failed: %v", err)
	}
	cleanedText := reg.ReplaceAllString(text, " ")
	return strings.Fields(strings.ToLower(cleanedText))
}

// BuildVocab создаёт словари для входных и выходных токенов
func BuildVocab(data []ParsedData) (map[string]int, map[string]int) {
	inputWordFreq := make(map[string]int)
	outputWordFreq := make(map[string]int)

	for _, item := range data {
		inputTokens := Tokenize(item.Input)
		outputTokens := Tokenize(item.Output)

		for _, token := range inputTokens {
			inputWordFreq[token]++
		}
		for _, token := range outputTokens {
			outputWordFreq[token]++
		}
	}

	buildSingleVocab := func(wordFreq map[string]int) map[string]int {
		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-1 { // -1 для <UNK>
			freqs = freqs[:maxVocabSize-1]
		}

		vocab := make(map[string]int)
		vocab["<UNK>"] = 0 // Специальный токен для неизвестных слов
		for i, kv := range freqs {
			vocab[kv.Key] = i + 1
		}
		return vocab
	}

	inputVocab := buildSingleVocab(inputWordFreq)
	outputVocab := buildSingleVocab(outputWordFreq)

	return inputVocab, outputVocab
}

// Vectorize превращает токены в индексы согласно словарю
func Vectorize(tokens []string, vocab map[string]int) []int {
	indices := make([]int, len(tokens))
	for i, token := range tokens {
		if idx, exists := vocab[token]; exists {
			indices[i] = idx
		} else {
			indices[i] = vocab["<UNK>"]
		}
	}
	return indices
}

// SaveWeights сохраняет матрицу весов в файл
func SaveWeights(weights *tensor.Dense, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

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

// LoadWeights загружает матрицу весов из файла
func LoadWeights(filename string) (*tensor.Dense, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

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

// 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
}

// EmbeddingLookup выполняет операцию поиска эмбеддингов с помощью one-hot кодирования
func EmbeddingLookup(g *gorgonia.ExprGraph, xIndices *gorgonia.Node, embeddings *gorgonia.Node, vocabSize int) *gorgonia.Node {
	// Создание one-hot кодирования
	oneHot, err := gorgonia.OneHot(xIndices, vocabSize)
	if err != nil {
		log.Fatalf("Ошибка при создании OneHot: %v", err)
	}

	// Умножение one-hot на матрицу эмбеддингов
	embedded, err := gorgonia.Mul(oneHot, embeddings)
	if err != nil {
		log.Fatalf("Ошибка при умножении OneHot на эмбеддинги: %v", err)
	}

	return gorgonia.Must(embedded) // Shape: [batchSize, maxSeqLength, embeddingSize]
}

// trainModel выполняет обучение модели
func trainModel(trainDir string, maxTrainMessages int) {
	parsedData, err := LoadCSVFiles(trainDir)
	if err != nil {
		log.Fatalf("Error loading CSV files: %v", err)
	}

	if maxTrainMessages > len(parsedData) {
		maxTrainMessages = len(parsedData)
	}

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

	rand.Seed(time.Now().UnixNano())

	inputVocab, outputVocab := BuildVocab(parsedData)
	fmt.Printf("Размер входного словаря: %d\n", len(inputVocab))
	fmt.Printf("Количество классов на выходе: %d\n", len(outputVocab))

	// Сохранение словарей
	if err := SaveVocabulary(inputVocab, "input_vocab.gob"); err != nil {
		log.Fatalf("Error saving input vocabulary: %v", err)
	}
	if err := SaveVocabulary(outputVocab, "output_vocab.gob"); err != nil {
		log.Fatalf("Error saving output vocabulary: %v", err)
	}

	totalSamples := len(parsedData)
	numClasses := len(outputVocab)

	g := gorgonia.NewGraph()

	// Инициализация эмбеддингов
	embeddings := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(len(inputVocab), embeddingSize), gorgonia.WithName("embeddings"), gorgonia.WithInit(gorgonia.GlorotU(1.0)))

	// Матрица весов для классификации
	W := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(embeddingSize, numClasses), gorgonia.WithName("W"), gorgonia.WithInit(gorgonia.GlorotN(1.0)))

	// Входные данные: индексы слов
	xIndices := gorgonia.NewMatrix(g, tensor.Int, gorgonia.WithShape(batchSize, maxSeqLength), gorgonia.WithName("xIndices"))

	// Метки: one-hot кодирование
	y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, numClasses), gorgonia.WithName("y"))

	// Операция поиска эмбеддингов
	embeddedNodes := EmbeddingLookup(g, xIndices, embeddings, len(inputVocab)) // Shape: [batchSize, maxSeqLength, embeddingSize]

	// Суммируем эмбеддинги по длине последовательности
	sumEmbeddings := gorgonia.Must(gorgonia.Sum(embeddedNodes, 1)) // Shape: [batchSize, embeddingSize]

	// Прогнозирование: yhat = sumEmbeddings * W
	yhat := gorgonia.Must(gorgonia.Mul(sumEmbeddings, W)) // Shape: [batchSize, numClasses]

	// Применение softmax
	yHatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat)) // Shape: [batchSize, numClasses]

	// Реализация кросс-энтропии
	logProbs := gorgonia.Must(gorgonia.Log(yHatSoftmax))
	crossEntropy := gorgonia.Must(gorgonia.Mul(y, logProbs))                           // [batchSize, numClasses]
	sumCrossEntropy := gorgonia.Must(gorgonia.Sum(crossEntropy, 1))                    // [batchSize]
	loss := gorgonia.Must(gorgonia.Mean(gorgonia.Must(gorgonia.Neg(sumCrossEntropy)))) // Scalar

	// Вычисление градиентов
	_, err = gorgonia.Grad(loss, W, embeddings)
	if err != nil {
		log.Fatalf("Не удалось создать градиенты: %v", err)
	}

	// Создание машины для выполнения графа
	machine := gorgonia.NewTapeMachine(g, gorgonia.BindDualValues(W, embeddings))
	defer machine.Close()

	optimizer := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(learningRate))

	for epoch := 0; epoch < numEpochs; epoch++ {
		// Перемешивание данных для каждой эпохи
		shuffledData := make([]ParsedData, len(parsedData))
		copy(shuffledData, parsedData)
		rand.Shuffle(len(shuffledData), func(i, j int) {
			shuffledData[i], shuffledData[j] = shuffledData[j], shuffledData[i]
		})

		totalLoss := 0.0
		for start := 0; start < totalSamples; start += batchSize {
			end := start + batchSize
			if end > totalSamples {
				end = totalSamples
			}

			currentBatchSize := end - start
			if currentBatchSize != batchSize {
				continue // Пропуск неполных батчей
			}

			// Подготовка данных батча
			xBatch := make([]int, batchSize*maxSeqLength)
			yBatch := make([]float32, batchSize*numClasses)

			for i := 0; i < batchSize; i++ {
				parsed := shuffledData[start+i]
				tokens := Tokenize(parsed.Input)
				indices := Vectorize(tokens, inputVocab)

				// Ограничение длины последовательности
				if len(indices) > maxSeqLength {
					indices = indices[:maxSeqLength]
				} else {
					// Дополнение последовательности до maxSeqLength
					for len(indices) < maxSeqLength {
						indices = append(indices, inputVocab["<UNK>"])
					}
				}

				for j, idx := range indices {
					xBatch[i*maxSeqLength+j] = idx
				}

				// One-hot кодирование выхода (используем прямое индексирование)
				outputClass, exists := outputVocab[parsed.Output]
				if !exists {
					outputClass = outputVocab["<UNK>"]
				}
				yBatch[i*numClasses+outputClass] = 1.0
			}

			xTensor := tensor.New(tensor.Of(tensor.Int), tensor.WithShape(batchSize, maxSeqLength), tensor.WithBacking(xBatch))
			yTensor := tensor.New(tensor.Of(tensor.Float32), tensor.WithShape(batchSize, numClasses), tensor.WithBacking(yBatch))

			// Обновление значений узлов
			gorgonia.Let(xIndices, xTensor)
			gorgonia.Let(y, yTensor)

			// Запуск машины
			if err := machine.RunAll(); err != nil {
				log.Fatalf("Ошибка выполнения машины: %v", err)
			}

			// Получение значения потерь
			lossVal := loss.Value().Data().(float32)
			totalLoss += float64(lossVal)

			// Обновление весов
			if err := optimizer.Step(gorgonia.NodesToValueGrads([]*gorgonia.Node{W, embeddings})); err != nil {
				log.Fatalf("Ошибка шага оптимизатора: %v", err)
			}

			// Сброс машины для следующей итерации
			machine.Reset()

			// Логирование каждых 100 батчей
			if (start/batchSize)%100 == 0 && start != 0 {
				avgLoss := totalLoss / 100.0
				fmt.Printf("Эпоха %d, Батч %d: Средняя потеря: %.4f\n", epoch+1, (start/batchSize)+1, avgLoss)
				totalLoss = 0.0
			}
		}
		fmt.Printf("Эпоха %d завершена.\n", epoch+1)
	}
	fmt.Println("Обучение модели завершено")

	// Сохранение весов модели
	if err := SaveWeights(W.Value().(*tensor.Dense), "model_weights.gob"); err != nil {
		log.Fatalf("Ошибка сохранения весов модели: %v", err)
	} else {
		fmt.Println("Веса модели сохранены успешно.")
	}

	// Сохранение эмбеддингов
	if err := SaveWeights(embeddings.Value().(*tensor.Dense), "embeddings.gob"); err != nil {
		log.Fatalf("Ошибка сохранения эмбеддингов: %v", err)
	} else {
		fmt.Println("Эмбеддинги сохранены успешно.")
	}
}

// argMax возвращает индекс максимального значения в слайсе
func argMax(values []float32) int {
	maxIdx := 0
	maxVal := values[0]
	for i, v := range values {
		if v > maxVal {
			maxVal = v
			maxIdx = i
		}
	}
	return maxIdx
}

// interactWithModel выполняет взаимодействие с моделью в режиме диалога
func interactWithModel() {
	W, err := LoadWeights("model_weights.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки весов модели: %v", err)
	}

	embeddings, err := LoadWeights("embeddings.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки эмбеддингов: %v", err)
	}

	inputVocab, err := LoadVocabulary("input_vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки входного словаря: %v", err)
	}
	outputVocab, err := LoadVocabulary("output_vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки выходного словаря: %v", err)
	}

	// Создание обратного словаря для выходных классов
	reverseOutputVocab := make(map[int]string)
	for word, idx := range outputVocab {
		reverseOutputVocab[idx] = word
	}

	g := gorgonia.NewGraph()
	embeddingsNode := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(len(inputVocab), embeddingSize), gorgonia.WithName("embeddings"), gorgonia.WithValue(embeddings))
	WNode := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(embeddingSize, len(outputVocab)), gorgonia.WithName("W"), gorgonia.WithValue(W))

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("Задайте вопрос (или введите 'exit' для выхода): ")
		text, _ := reader.ReadString('\n')
		text = strings.TrimSpace(text)

		if text == "exit" {
			break
		}

		tokens := Tokenize(text)
		indices := Vectorize(tokens, inputVocab)

		// Ограничение длины последовательности
		if len(indices) > maxSeqLength {
			indices = indices[:maxSeqLength]
		} else {
			// Дополнение последовательности до maxSeqLength
			for len(indices) < maxSeqLength {
				indices = append(indices, inputVocab["<UNK>"])
			}
		}

		// Создание батча из одного примера
		xBatch := make([]int, maxSeqLength)
		copy(xBatch, indices)
		xTensor := tensor.New(tensor.Of(tensor.Int), tensor.WithShape(1, maxSeqLength), tensor.WithBacking(xBatch))

		// Входные индексы
		xIndicesNode := gorgonia.NewMatrix(g, tensor.Int, gorgonia.WithShape(1, maxSeqLength), gorgonia.WithName("xIndices"))

		// Обновление значений узла xIndices
		gorgonia.Let(xIndicesNode, xTensor)

		// Операция поиска эмбеддингов
		embeddedNodes := EmbeddingLookup(g, xIndicesNode, embeddingsNode, len(inputVocab)) // Shape: [1, maxSeqLength, embeddingSize]

		// Суммируем эмбеддинги по длине последовательности
		sumEmbeddings := gorgonia.Must(gorgonia.Sum(embeddedNodes, 1)) // Shape: [1, embeddingSize]

		// Прогнозирование
		yhat := gorgonia.Must(gorgonia.Mul(sumEmbeddings, WNode)) // Shape: [1, numClasses]

		// Применение softmax
		yHatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat)) // Shape: [1, numClasses]

		// Создание машины для выполнения графа
		machine := gorgonia.NewTapeMachine(g)
		if err := machine.RunAll(); err != nil {
			log.Fatalf("Ошибка выполнения машины: %v", err)
		}

		predictionsVal := yHatSoftmax.Value().Data().([]float32)
		predictedClass := argMax(predictionsVal)

		response, exists := reverseOutputVocab[predictedClass]
		if !exists {
			response = "Извините, я не понимаю."
		}

		fmt.Printf("Ответ: %s\n\n", response)
		machine.Reset()
	}
}

// validateModel выполняет валидацию модели
func validateModel(validationDir string) {
	W, err := LoadWeights("model_weights.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки весов модели: %v", err)
	}

	embeddings, err := LoadWeights("embeddings.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки эмбеддингов: %v", err)
	}

	inputVocab, err := LoadVocabulary("input_vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки входного словаря: %v", err)
	}
	outputVocab, err := LoadVocabulary("output_vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки выходного словаря: %v", err)
	}

	validationData, err := LoadCSVFiles(validationDir)
	if err != nil {
		log.Fatalf("Ошибка загрузки CSV файлов для валидации: %v", err)
	}

	g := gorgonia.NewGraph()
	embeddingsNode := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(len(inputVocab), embeddingSize), gorgonia.WithName("embeddings"), gorgonia.WithValue(embeddings))
	WNode := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(embeddingSize, len(outputVocab)), gorgonia.WithName("W"), gorgonia.WithValue(W))

	correct := 0
	total := len(validationData)

	for _, item := range validationData {
		inputTokens := Tokenize(item.Input)
		indices := Vectorize(inputTokens, inputVocab)

		// Ограничение длины последовательности
		if len(indices) > maxSeqLength {
			indices = indices[:maxSeqLength]
		} else {
			for len(indices) < maxSeqLength {
				indices = append(indices, inputVocab["<UNK>"])
			}
		}

		// Создание батча из одного примера
		xBatch := make([]int, maxSeqLength)
		copy(xBatch, indices)
		xTensor := tensor.New(tensor.Of(tensor.Int), tensor.WithShape(1, maxSeqLength), tensor.WithBacking(xBatch))

		// Входные индексы
		xIndicesNode := gorgonia.NewMatrix(g, tensor.Int, gorgonia.WithShape(1, maxSeqLength), gorgonia.WithName("xIndices"))

		// Обновление значений узла xIndices
		gorgonia.Let(xIndicesNode, xTensor)

		// Операция поиска эмбеддингов
		embeddedNodes := EmbeddingLookup(g, xIndicesNode, embeddingsNode, len(inputVocab)) // Shape: [1, maxSeqLength, embeddingSize]

		// Суммируем эмбеддинги по длине последовательности
		sumEmbeddings := gorgonia.Must(gorgonia.Sum(embeddedNodes, 1)) // Shape: [1, embeddingSize]

		// Прогнозирование
		yhat := gorgonia.Must(gorgonia.Mul(sumEmbeddings, WNode)) // Shape: [1, numClasses]

		// Применение softmax
		yHatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat)) // Shape: [1, numClasses]

		// Создание машины для выполнения графа
		machine := gorgonia.NewTapeMachine(g)
		if err := machine.RunAll(); err != nil {
			log.Fatalf("Ошибка выполнения машины: %v", err)
		}

		predictionsVal := yHatSoftmax.Value().Data().([]float32)
		predictedClass := argMax(predictionsVal)

		// Определение истинного класса
		trueClass, exists := outputVocab[item.Output]
		if !exists {
			trueClass = outputVocab["<UNK>"]
		}

		if predictedClass == trueClass {
			correct++
		}

		machine.Reset()
	}

	accuracy := float64(correct) / float64(total) * 100
	fmt.Printf("Валидационная точность: %.2f%%\n", accuracy)
}

func main() {
	fmt.Print("Введите режим работы (t для обучения, v для валидации, d для диалога): ")
	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("Неизвестный режим. Используйте 't' для обучения, 'v' для валидации или 'd' для диалога.")
	}
}


Not enough arguments in call to 'gorgonia.Must' ощибка Unresolved reference 'OneHot' не существует