package main

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

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

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

// Константы для ограничения количества сообщений и размера словаря
const (
	maxTrainMessages = 9900
	maxVocabSize     = 5000000
	numEpochs        = 10
	batchSize        = 150
)

// 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") || strings.HasSuffix(d.Name(), ".dt") {
			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.TrimSpace(text))
}

// BuildVocab строит словарь на основе данных и возвращает отображение слова в индекс
func BuildVocab(data []ParsedData) map[string]int {
	vocab := make(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]
	}

	for i, kv := range freqs {
		vocab[kv.Key] = i
	}
	return vocab
}

// Vectorize преобразует токены в вектор на основе словаря
func Vectorize(tokens []string, vocab map[string]int) []float32 {
	vector := make([]float32, len(vocab))
	for _, token := range tokens {
		if idx, exists := vocab[token]; exists {
			vector[idx] += 1
		}
	}
	return vector
}

// Сохранение и загрузка весов модели
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)
}

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
}

// Сохранение и загрузка словаря
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)
}

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
}

// Сохранение и загрузка сопоставления классов с ответами
func SaveResponseMap(classToOutput map[int]string, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

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

func LoadResponseMap(filename string) (map[int]string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

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

// Сохранение и загрузка сопоставления выходных строк к классам
func SaveOutputToClass(outputToClass 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(outputToClass)
}

func LoadOutputToClass(filename string) (map[string]int, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

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

// BuildOutputClasses создает сопоставление между выходными строками и классами
func BuildOutputClasses(data []ParsedData) (map[string]int, map[int]string) {
	outputToClass := make(map[string]int)
	classToOutput := make(map[int]string)
	classIndex := 0

	for _, item := range data {
		if _, exists := outputToClass[item.Output]; !exists {
			outputToClass[item.Output] = classIndex
			classToOutput[classIndex] = item.Output
			classIndex++
		}
	}

	return outputToClass, classToOutput
}

// TrainModel обучает модель многоуровневой классификации
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)
	outputToClass, classToOutput := BuildOutputClasses(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)
	}

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

	if err := SaveOutputToClass(outputToClass, "models/output_to_class.gob"); err != nil {
		log.Fatalf("Ошибка при сохранении output_to_class: %v", err)
	}

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

	g := gorgonia.NewGraph()

	inputSize := len(vocab)
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, numClasses), gorgonia.WithInit(gorgonia.GlorotN(1.0)), gorgonia.WithName("weights"))

	x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, inputSize), gorgonia.WithName("x"))
	y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, numClasses), gorgonia.WithName("y"))

	// Предсказания
	yhat := gorgonia.Must(gorgonia.Mul(x, w))
	yhatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat))

	// Функция потерь: кросс-энтропия
	// Реализуем вручную: -sum(y * log(yhat))
	logYhat := gorgonia.Must(gorgonia.Log(yhatSoftmax))
	prod := gorgonia.Must(gorgonia.HadamardProd(y, logYhat))
	sum := gorgonia.Must(gorgonia.Sum(prod, 1)) // сумма по классам
	loss := gorgonia.Must(gorgonia.Neg(sum))
	averageLoss := gorgonia.Must(gorgonia.Mean(loss))

	// Создаем градиенты
	if _, err := gorgonia.Grad(averageLoss, w); err != nil {
		log.Fatalf("Не удалось создать градиенты: %v", err)
	}

	// Инициализируем оптимизатор
	optimizer := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(0.001))

	// Создаем машину
	machine := gorgonia.NewTapeMachine(g, gorgonia.BindDualValues(w))
	defer machine.Close()

	for epoch := 0; epoch < numEpochs; epoch++ {
		var epochLossTotal float64
		barWidth := 50 // Ширина прогресс-бара

		// Инициализируем прогресс-бар
		fmt.Printf("Эпоха %d/%d\n[", epoch+1, numEpochs)
		for i := 0; i < barWidth; i++ {
			fmt.Print(" ")
		}
		fmt.Print("] 0%")
		fmt.Print("\r[") // Перемещаем курсор в начало прогресс-бара

		for start := 0; start < totalSamples; start += batchSize {
			end := start + batchSize
			batchSizeActual := batchSize

			if end > totalSamples {
				batchSizeActual = totalSamples - start
				end = totalSamples
			}

			// Инициализируем данные с паддингом
			xValData := make([]float32, batchSizeActual*inputSize)  // По умолчанию заполнено нулями
			yValData := make([]float32, batchSizeActual*numClasses) // One-hot векторы

			for i := start; i < end; i++ {
				tokens := Tokenize(parsedData[i].Input)
				inputVec := Vectorize(tokens, vocab)
				idx := i - start

				copy(xValData[idx*inputSize:(idx+1)*inputSize], inputVec)

				classIdx, exists := outputToClass[parsedData[i].Output]
				if !exists {
					continue // Пропускаем, если класс не найден
				}
				yValData[idx*numClasses+classIdx] = 1.0
			}

			xVal := tensor.New(tensor.WithShape(batchSizeActual, inputSize), tensor.Of(tensor.Float32), tensor.WithBacking(xValData))
			yVal := tensor.New(tensor.WithShape(batchSizeActual, numClasses), tensor.Of(tensor.Float32), tensor.WithBacking(yValData))

			// Устанавливаем значения узлов с помощью gorgonia.Let
			if err := gorgonia.Let(x, xVal); err != nil {
				log.Fatalf("Ошибка при установке значения x: %v", err)
			}
			if err := gorgonia.Let(y, yVal); err != nil {
				log.Fatalf("Ошибка при установке значения y: %v", err)
			}

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

			// Получаем значение функции потерь
			lossValInterface := averageLoss.Value()
			if lossValInterface == nil {
				log.Fatalf("Потери пусты")
			}
			lossValFloat := lossValInterface.Data().([]float32)[0]
			epochLossTotal += float64(lossValFloat) * float64(batchSizeActual)

			// Обновляем веса
			if err := optimizer.Step(w); err != nil {
				log.Fatalf("Ошибка при обновлении весов: %v", err)
			}

			// Обновляем прогресс-бар
			currentBatch := start + batchSizeActual
			if currentBatch > totalSamples {
				currentBatch = totalSamples
			}
			percent := float64(currentBatch) / float64(totalSamples)
			filledLength := int(float64(barWidth) * percent) // Преобразуем barWidth в float64 перед умножением

			// Перемещаем курсор в начало прогресс-бара
			fmt.Printf("\r[") // Сбрасываем на начало строки прогресс-бара
			for i := 0; i < filledLength; i++ {
				fmt.Print("=")
			}
			for i := filledLength; i < barWidth; i++ {
				fmt.Print(" ")
			}
			fmt.Printf("] %3.0f%%", percent*100)

			machine.Reset()
		}

		averageLossEpoch := epochLossTotal / float64(totalSamples)
		// Переходим на новую строку после завершения прогресс-бара
		fmt.Printf("\r[")
		for i := 0; i < barWidth; i++ {
			fmt.Print("=")
		}
		fmt.Printf("] 100%% - Средняя Потеря: %.4f\n", averageLossEpoch)
	}

	fmt.Println("Обучение модели завершено")

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

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

	vocab, err := LoadVocabulary("models/vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке словаря: %v", err)
	}

	classToOutput, err := LoadResponseMap("models/response_map.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке сопоставления классов: %v", err)
	}

	inputSize := len(vocab)
	numClasses := len(classToOutput)
	reader := bufio.NewReader(os.Stdin)

	// Создаем граф для предсказания
	g := gorgonia.NewGraph()
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, numClasses), gorgonia.WithValue(weightsTensor), gorgonia.WithName("weights"))

	x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, inputSize), gorgonia.WithName("x"))
	yhat := gorgonia.Must(gorgonia.Mul(x, w))
	yhatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat))

	machine := gorgonia.NewTapeMachine(g, gorgonia.BindDualValues(w))
	defer machine.Close()

	for {
		fmt.Print("Вы: ")
		text, _ := reader.ReadString('\n')
		text = strings.TrimSpace(text)

		if strings.ToLower(text) == "exit" {
			fmt.Println("До свидания!")
			break
		}

		tokens := Tokenize(text)
		inputVec := Vectorize(tokens, vocab)

		xValData := make([]float32, inputSize)
		copy(xValData, inputVec)
		xVal := tensor.New(tensor.WithShape(1, inputSize), tensor.Of(tensor.Float32), tensor.WithBacking(xValData))

		// Устанавливаем значение узла x с помощью gorgonia.Let
		if err := gorgonia.Let(x, xVal); err != nil {
			log.Fatalf("Ошибка при установке значения x: %v", err)
		}

		if err := machine.RunAll(); err != nil {
			log.Fatalf("Ошибка во время взаимодействия: %v", err)
		}

		// Получаем предсказания
		if yhatSoftmax.Value() == nil {
			log.Fatalf("Predicted value is nil")
		}
		predictionsVal := yhatSoftmax.Value().Data().([]float32)
		predictedClass := 0
		maxProb := predictionsVal[0]
		for i, prob := range predictionsVal {
			if prob > maxProb {
				maxProb = prob
				predictedClass = i
			}
		}

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

		fmt.Printf("ИИ: %s (Уверенность: %.2f%%)\n", response, maxProb*100)
		machine.Reset()
	}
}

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

	vocab, err := LoadVocabulary("models/vocab.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке словаря: %v", err)
	}

	classToOutput, err := LoadResponseMap("models/response_map.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке сопоставления классов: %v", err)
	}

	inputSize := len(vocab)
	numClasses := len(classToOutput)

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

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

	g := gorgonia.NewGraph()
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, numClasses), gorgonia.WithValue(weightsTensor), gorgonia.WithName("weights"))

	x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, inputSize), gorgonia.WithName("x"))
	y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, numClasses), gorgonia.WithName("y"))

	yhat := gorgonia.Must(gorgonia.Mul(x, w))
	yhatSoftmax := gorgonia.Must(gorgonia.SoftMax(yhat))

	// Валидация не требует градиентов
	machine := gorgonia.NewTapeMachine(g, gorgonia.BindDualValues(w))
	defer machine.Close()

	var totalLoss float64
	correct := 0
	total := len(validationData)

	for idx, item := range validationData {
		inputTokens := Tokenize(item.Input)
		inputVec := Vectorize(inputTokens, vocab)

		xValData := make([]float32, inputSize)
		copy(xValData, inputVec)
		xVal := tensor.New(tensor.WithShape(1, inputSize), tensor.Of(tensor.Float32), tensor.WithBacking(xValData))

		// Получаем индекс класса для выхода
		classIdx, exists := classToOutputReverse(item.Output, classToOutput)
		if !exists {
			// Если выходной класс не был виден во время обучения, пропускаем пример
			continue
		}
		yValData := make([]float32, numClasses)
		yValData[classIdx] = 1.0
		yVal := tensor.New(tensor.WithShape(1, numClasses), tensor.Of(tensor.Float32), tensor.WithBacking(yValData))

		// Устанавливаем значения узлов с помощью gorgonia.Let
		if err := gorgonia.Let(x, xVal); err != nil {
			log.Fatalf("Ошибка при установке значения x: %v", err)
		}
		if err := gorgonia.Let(y, yVal); err != nil {
			log.Fatalf("Ошибка при установке значения y: %v", err)
		}

		if err := machine.RunAll(); err != nil {
			log.Fatalf("Ошибка во время валидации: %v", err)
		}

		// Получаем предсказания
		if yhatSoftmax.Value() == nil {
			log.Fatalf("Predicted value is nil")
		}
		predictionsVal := yhatSoftmax.Value().Data().([]float32)
		predictedClass := 0
		maxProb := predictionsVal[0]
		for i, prob := range predictionsVal {
			if prob > maxProb {
				maxProb = prob
				predictedClass = i
			}
		}

		expectedClass := classIdx
		if predictedClass == expectedClass {
			correct++
		}

		// Вычисляем кросс-энтропию для текущего примера
		epsilon := float32(1e-7) // Для предотвращения log(0)
		p := float32(math.Min(math.Max(float64(predictionsVal[predictedClass]), float64(epsilon)), 1-1e-7))
		lossVal := -yValData[predictedClass]*float32(math.Log(float64(p))) -
			(1-yValData[predictedClass])*float32(math.Log(float64(1-p)))
		totalLoss += float64(lossVal)

		machine.Reset()

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

	averageLoss := totalLoss / float64(total)
	accuracy := float64(correct) / float64(total) * 100
	fmt.Printf("Потери на валидации: %.4f\n", averageLoss)
	fmt.Printf("Точность на валидации: %.2f%%\n", accuracy)
}

// classToOutputReverse помогает найти индекс класса по выходной строке
func classToOutputReverse(output string, classToOutput map[int]string) (int, bool) {
	for k, v := range classToOutput {
		if v == output {
			return k, true
		}
	}
	return -1, false
}

// main - точка входа программы
func main() {
	fmt.Print("Введите режим работы (t для обучения, v для валидации, d для диалога): ")
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	mode := scanner.Text()

	switch mode {
	case "t":
		fmt.Print("Введите директорию для обучения: ")
		scanner.Scan()
		trainDir := scanner.Text()
		TrainModel(trainDir, maxTrainMessages)
	case "v":
		fmt.Print("Введите директорию для валидации: ")
		scanner.Scan()
		validationDir := scanner.Text()
		ValidateModel(validationDir)
	case "d":
		InteractWithModel()
	default:
		fmt.Println("Неизвестный режим. Используйте 't' для обучения, 'v' для валидации или 'd' для диалога.")
	}
}

Cannot use 'w' (type *Node) as the type []ValueGrad


			// Обновляем веса
			if err := optimizer.Step(w); err != nil {
				log.Fatalf("Ошибка при обновлении весов: %v", err)
			}