package main

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

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

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

const (
	maxTrainMessages = 1000
	maxVocabSize     = 500000
	maxOutputClasses = 10000 // Максимальное количество уникальных выходных классов
	epochs           = 10
	learningRate     = 0.001
	batchSize        = 150
)

// Vocab хранит словари для входных и выходных данных
type Vocab struct {
	InputVocab  map[string]int
	OutputVocab map[string]int
}

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

// BuildVocabularies строит словари для входных и выходных данных
func BuildVocabularies(data []ParsedData) Vocab {
	inputVocab := make(map[string]int)
	outputVocab := make(map[string]int)
	wordFreq := make(map[string]int)

	// Подсчет частоты слов только для входных данных
	for _, item := range data {
		tokens := Tokenize(item.Input)
		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 {
		inputVocab[kv.Key] = i
	}

	// Построение выходного вокабуляра
	for _, item := range data {
		output := strings.ToLower(strings.TrimSpace(item.Output))
		if _, exists := outputVocab[output]; !exists && len(outputVocab) < maxOutputClasses {
			outputVocab[output] = len(outputVocab)
		}
	}

	return Vocab{
		InputVocab:  inputVocab,
		OutputVocab: outputVocab,
	}
}

// VectorizeInput преобразует токены в вектор признаков
func VectorizeInput(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, bias *tensor.Dense, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	modelData := struct {
		Weights *tensor.Dense
		Bias    *tensor.Dense
	}{
		Weights: weights,
		Bias:    bias,
	}

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

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

	var modelData struct {
		Weights *tensor.Dense
		Bias    *tensor.Dense
	}
	decoder := gob.NewDecoder(f)
	if err := decoder.Decode(&modelData); err != nil {
		return nil, nil, err
	}
	return modelData.Weights, modelData.Bias, nil
}

// Сохранение словарей
func SaveVocabulary(vocab Vocab, 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) (Vocab, error) {
	f, err := os.Open(filename)
	if err != nil {
		return Vocab{}, err
	}
	defer f.Close()

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

// 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("Loaded %d training records\n", len(parsedData))

	vocab := BuildVocabularies(parsedData)
	inputVocab := vocab.InputVocab
	outputVocab := vocab.OutputVocab
	numOutputClasses := len(outputVocab)
	fmt.Printf("Input vocabulary size: %d\n", len(inputVocab))
	fmt.Printf("Output classes: %d\n", numOutputClasses)

	if err := SaveVocabulary(vocab, "vocab.gob"); err != nil {
		log.Fatalf("Error saving vocabulary: %v", err)
	}

	totalSamples := len(parsedData)

	// Создание графа
	g := gorgonia.NewGraph()

	// Размеры входа и выхода
	inputSize := len(inputVocab)
	outputSize := numOutputClasses

	// Параметры модели (Весовая матрица и смещение)
	weights := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, outputSize), gorgonia.WithInit(gorgonia.GlorotU(1.0)), gorgonia.WithName("weights"))
	// Изменение формы bias на (1, outputSize)
	bias := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, outputSize), gorgonia.WithInit(gorgonia.Zeroes()), gorgonia.WithName("bias"))

	// Определение параметров для оптимизатора
	params := []*gorgonia.Node{weights, bias}

	// Определение обучающего оптимизатора
	adam := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(learningRate))

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

	for epoch := 0; epoch < epochs; epoch++ {
		fmt.Printf("Starting epoch %d/%d\n", epoch+1, epochs)
		for start := 0; start < totalSamples; start += batchSize {
			end := start + batchSize
			if end > totalSamples {
				end = totalSamples
			}

			batchData := parsedData[start:end]
			batchSz := end - start

			// Подготовка входных данных и меток
			xData := make([]float32, batchSz*inputSize)
			yData := make([]float32, batchSz*numOutputClasses)

			for i, item := range batchData {
				tokens := Tokenize(item.Input)
				inputVec := VectorizeInput(tokens, inputVocab)
				copy(xData[i*inputSize:(i+1)*inputSize], inputVec)

				// One-hot кодирование меток
				outputIdx, exists := outputVocab[strings.ToLower(strings.TrimSpace(item.Output))]
				if !exists {
					// Если метка не была встречена при построении вокабуляра, пропускаем запись
					continue
				}
				yData[i*numOutputClasses+outputIdx] = 1
			}

			// Создание тензоров
			xTensor := tensor.New(tensor.WithShape(batchSz, inputSize), tensor.Of(tensor.Float32), tensor.WithBacking(xData))
			yTensor := tensor.New(tensor.WithShape(batchSz, numOutputClasses), tensor.Of(tensor.Float32), tensor.WithBacking(yData))

			// Узлы ввода
			x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSz, inputSize), gorgonia.WithValue(xTensor), gorgonia.WithName("x"))
			y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSz, numOutputClasses), gorgonia.WithValue(yTensor), gorgonia.WithName("y"))

			// Прогноз
			logits := gorgonia.Must(gorgonia.Add(gorgonia.Must(gorgonia.Mul(x, weights)), bias))

			// Реализация функции потерь (кросс-энтропия)
			probs := gorgonia.Must(gorgonia.SoftMax(logits))
			logProbs := gorgonia.Must(gorgonia.Log(probs))
			crossEntropy := gorgonia.Must(gorgonia.Neg(
				gorgonia.Must(gorgonia.Mean(
					gorgonia.Must(gorgonia.Sum(
						gorgonia.Must(gorgonia.Mul(y, logProbs)),
						1,
					)),
				)),
			))
			loss := crossEntropy

			// Расчет градиентов
			_, err := gorgonia.Grad(loss, params...)
			if err != nil {
				log.Fatalf("Failed to calculate gradients: %v", err)
			}

			// Запуск вычислений
			if err := machine.RunAll(); err != nil {
				log.Fatalf("Machine run error: %v", err)
			}

			// Шаг оптимизатора
			adam.Step(gorgonia.NodesToValueGrads(params))

			// Сброс градиентов и машины
			machine.Reset()

			// Отображение прогресса
			progress := float64(end) / float64(totalSamples) * 100
			fmt.Printf("Epoch %d: Processed batch from %d to %d (%.2f%% complete)\n", epoch+1, start, end, progress)
		}
	}

	fmt.Println("Model training complete")

	// Сохранение весов модели и смещений
	if err := SaveWeights(weights.Value().(*tensor.Dense), bias.Value().(*tensor.Dense), "model_weights.gob"); err != nil {
		log.Fatalf("Error saving model weights: %v", err)
	} else {
		fmt.Println("Model weights saved successfully.")
	}
}

// interactWithModel позволяет взаимодействовать с обученной моделью
func interactWithModel() {
	// Загрузка весов и смещений
	weightsTensor, biasTensor, err := LoadWeights("model_weights.gob")
	if err != nil {
		log.Fatalf("Error loading model weights: %v", err)
	}

	// Загрузка словарей
	vocab, err := LoadVocabulary("vocab.gob")
	if err != nil {
		log.Fatalf("Error loading vocabulary: %v", err)
	}

	inputVocab := vocab.InputVocab
	outputVocab := vocab.OutputVocab

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

	inputSize := len(inputVocab)
	outputSize := len(outputVocab)

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

	reader := bufio.NewReader(os.Stdin)

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

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

		tokens := Tokenize(text)
		inputVec := VectorizeInput(tokens, inputVocab)

		xVal := tensor.New(tensor.WithShape(1, inputSize), tensor.Of(tensor.Float32), tensor.WithBacking(inputVec))
		x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, inputSize), gorgonia.WithValue(xVal), gorgonia.WithName("x"))

		// Прогноз
		logits := gorgonia.Must(gorgonia.Add(gorgonia.Must(gorgonia.Mul(x, w)), b))
		probs := gorgonia.Must(gorgonia.SoftMax(logits))

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

		if err := machine.RunAll(); err != nil {
			log.Fatalf("Machine run error: %v", err)
		}

		probsVal := probs.Value().Data().([]float32)

		// Нахождение индекса с максимальной вероятностью
		maxIdx := 0
		maxProb := probsVal[0]
		for i, prob := range probsVal {
			if prob > maxProb {
				maxProb = prob
				maxIdx = i
			}
		}

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

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

// validateModel валидирует модель на предоставленных данных
func validateModel(validationDir string) {
	// Загрузка весов и смещений
	weightsTensor, biasTensor, err := LoadWeights("model_weights.gob")
	if err != nil {
		log.Fatalf("Error loading model weights: %v", err)
	}

	// Загрузка словарей
	vocab, err := LoadVocabulary("vocab.gob")
	if err != nil {
		log.Fatalf("Error loading vocabulary: %v", err)
	}

	inputVocab := vocab.InputVocab
	outputVocab := vocab.OutputVocab
	numOutputClasses := len(outputVocab)

	validationData, err := LoadCSVFiles(validationDir)
	if err != nil {
		log.Fatalf("Error loading validation CSV files: %v", err)
	}

	fmt.Printf("Loaded %d validation records\n", len(validationData))

	total := 0
	correct := 0

	// Создание графа для валидации
	g := gorgonia.NewGraph()
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(len(inputVocab), numOutputClasses), gorgonia.WithValue(weightsTensor), gorgonia.WithName("weights"))
	// Изменение формы bias на (1, outputSize)
	b := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, numOutputClasses), gorgonia.WithValue(biasTensor), gorgonia.WithName("bias"))

	for _, item := range validationData {
		tokens := Tokenize(item.Input)
		inputVec := VectorizeInput(tokens, inputVocab)

		xVal := tensor.New(tensor.WithShape(1, len(inputVocab)), tensor.Of(tensor.Float32), tensor.WithBacking(inputVec))
		x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, len(inputVocab)), gorgonia.WithValue(xVal), gorgonia.WithName("x"))

		// Прогноз
		logits := gorgonia.Must(gorgonia.Add(gorgonia.Must(gorgonia.Mul(x, w)), b))
		probs := gorgonia.Must(gorgonia.SoftMax(logits))

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

		if err := machine.RunAll(); err != nil {
			log.Fatalf("Machine run error: %v", err)
		}

		probsVal := probs.Value().Data().([]float32)

		// Нахождение индекса с максимальной вероятностью
		maxIdx := 0
		maxProb := probsVal[0]
		for i, prob := range probsVal {
			if prob > maxProb {
				maxProb = prob
				maxIdx = i
			}
		}

		// Истинная метка
		trueOutput := strings.ToLower(strings.TrimSpace(item.Output))
		trueIdx, exists := outputVocab[trueOutput]
		if !exists {
			continue // Пропускаем, если метка не была встречена при обучении
		}

		if maxIdx == trueIdx {
			correct++
		}
		total++

		machine.Reset()
	}

	if total == 0 {
		fmt.Println("Нет валидируемых данных.")
		return
	}

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

// main является точкой входа в программу
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' для диалога.")
	}
}


исправь ошибку

Введите режим работы (t для обучения, v для валидации, d для диалога): t
Loaded 1000 training records
Input vocabulary size: 3297
Output classes: 680
Starting epoch 1/10
panic: Failed to infer shape. Op: + false: Shape mismatch: (150, 680) and (1, 680)

goroutine 1 [running]:
gorgonia.org/gorgonia.Must(...)
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/gorgonia.go:19
main.trainModel({0xed1309?, 0xc00009c020?}, 0x3e8)
        D:/JBProjects/GOLAND/Learn/main.go:304 +0x17d4
main.main()
        D:/JBProjects/GOLAND/Learn/main.go:533 +0x11e