package main

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

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

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

// Константы
const (
	maxTrainMessages = 10000
	maxVocabSize     = 50000
	numEpochs        = 10
	batchSize        = 64
	embeddingSize    = 128
	hiddenSize       = 256
	startToken       = "<SOS>"
	endToken         = "<EOS>"
	padToken         = "<PAD>"
	unkToken         = "<UNK>"
)

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

		if strings.HasSuffix(info.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(info.Name(), ".txt") || strings.HasSuffix(info.Name(), ".dt") {
			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer f.Close()

			reader := bufio.NewReader(f)
			for {
				line, err := reader.ReadString('\n')
				if err != nil {
					if err == io.EOF {
						break
					}
					return err
				}
				if len(data) >= maxMessages {
					break
				}
				line = strings.TrimSpace(line)
				if line == "" {
					continue
				}
				parts := strings.SplitN(line, ":", 2)
				if len(parts) == 2 {
					data = append(data, ParsedData{
						Input:  strings.TrimSpace(parts[0]),
						Output: strings.TrimSpace(parts[1]),
					})
				}
			}
		}
		return nil
	})
	return data, err
}

func Tokenize(text string) []string {
	return strings.Fields(strings.TrimSpace(strings.ToLower(text)))
}

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

	vocab := make(map[string]int)
	vocab[padToken] = 0
	vocab[startToken] = 1
	vocab[endToken] = 2
	vocab[unkToken] = 3

	idx := 4
	for _, kv := range freqs {
		if idx >= maxVocabSize {
			break
		}
		vocab[kv.Key] = idx
		idx++
	}
	return vocab
}

func Vectorize(tokens []string, vocab map[string]int) []int {
	var vector []int
	for _, token := range tokens {
		if idx, exists := vocab[token]; exists {
			vector = append(vector, idx)
		} else {
			vector = append(vector, vocab[unkToken])
		}
	}
	return vector
}

// Сохранение и загрузка словаря
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 SaveModel(params []*gorgonia.Node, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	values := make([]tensor.Tensor, len(params))
	for i, param := range params {
		values[i] = param.Value().(tensor.Tensor)
	}

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

func LoadModel(g *gorgonia.ExprGraph, shapes []tensor.Shape, filename string) ([]*gorgonia.Node, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var values []tensor.Tensor
	decoder := gob.NewDecoder(f)
	if err := decoder.Decode(&values); err != nil {
		return nil, err
	}

	if len(values) != len(shapes) {
		return nil, fmt.Errorf("number of tensors does not match shapes")
	}

	params := make([]*gorgonia.Node, len(shapes))
	for i, shape := range shapes {
		params[i] = gorgonia.NewTensor(g, tensor.Float32, shape.Dims(), gorgonia.WithShape(shape...), gorgonia.WithValue(values[i]))
	}

	return params, nil
}

// Модель Seq2Seq
type Seq2Seq struct {
	encoderEmbedding *gorgonia.Node
	encoderLSTM      *LSTM

	decoderEmbedding *gorgonia.Node
	decoderLSTM      *LSTM
	decoderDense     *gorgonia.Node

	vocabSize     int
	embeddingSize int
	hiddenSize    int
}

// LSTM слой
type LSTM struct {
	wf *gorgonia.Node
	uf *gorgonia.Node
	bf *gorgonia.Node

	wi *gorgonia.Node
	ui *gorgonia.Node
	bi *gorgonia.Node

	wo *gorgonia.Node
	uo *gorgonia.Node
	bo *gorgonia.Node

	wg *gorgonia.Node
	ug *gorgonia.Node
	bg *gorgonia.Node

	hiddenSize int
}

// Создаем новый LSTM слой
func NewLSTM(g *gorgonia.ExprGraph, inputSize, hiddenSize int) *LSTM {
	wf := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	uf := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(hiddenSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	bf := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))

	wi := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	ui := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(hiddenSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	bi := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))

	wo := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	uo := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(hiddenSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	bo := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))

	wg := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	ug := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(hiddenSize, hiddenSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))
	bg := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))

	return &LSTM{
		wf: wf, uf: uf, bf: bf,
		wi: wi, ui: ui, bi: bi,
		wo: wo, uo: uo, bo: bo,
		wg: wg, ug: ug, bg: bg,
		hiddenSize: hiddenSize,
	}
}

// Прямой проход через LSTM слой
func (lstm *LSTM) Forward(x, hPrev, cPrev *gorgonia.Node) (h, c *gorgonia.Node, err error) {
	// Forget gate
	ft := gorgonia.Must(gorgonia.Sigmoid(gorgonia.Must(
		gorgonia.Add(
			gorgonia.Must(gorgonia.Add(
				gorgonia.Must(gorgonia.Mul(x, lstm.wf)),
				gorgonia.Must(gorgonia.Mul(hPrev, lstm.uf)),
			)),
			lstm.bf,
		),
	)))

	// Input gate
	it := gorgonia.Must(gorgonia.Sigmoid(gorgonia.Must(
		gorgonia.Add(
			gorgonia.Must(gorgonia.Add(
				gorgonia.Must(gorgonia.Mul(x, lstm.wi)),
				gorgonia.Must(gorgonia.Mul(hPrev, lstm.ui)),
			)),
			lstm.bi,
		),
	)))

	// Candidate memory cell
	gt := gorgonia.Must(gorgonia.Tanh(gorgonia.Must(
		gorgonia.Add(
			gorgonia.Must(gorgonia.Add(
				gorgonia.Must(gorgonia.Mul(x, lstm.wg)),
				gorgonia.Must(gorgonia.Mul(hPrev, lstm.ug)),
			)),
			lstm.bg,
		),
	)))

	// Output gate
	ot := gorgonia.Must(gorgonia.Sigmoid(gorgonia.Must(
		gorgonia.Add(
			gorgonia.Must(gorgonia.Add(
				gorgonia.Must(gorgonia.Mul(x, lstm.wo)),
				gorgonia.Must(gorgonia.Mul(hPrev, lstm.uo)),
			)),
			lstm.bo,
		),
	)))

	// New memory cell
	c = gorgonia.Must(gorgonia.Add(gorgonia.Must(gorgonia.HadamardProd(ft, cPrev)), gorgonia.Must(gorgonia.HadamardProd(it, gt))))

	// New hidden state
	h = gorgonia.Must(gorgonia.HadamardProd(ot, gorgonia.Must(gorgonia.Tanh(c))))

	return h, c, nil
}

func NewSeq2Seq(g *gorgonia.ExprGraph, vocabSize, embeddingSize, hiddenSize int) *Seq2Seq {
	// Encoder embedding
	encoderEmbedding := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(vocabSize, embeddingSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))

	// Encoder LSTM
	encoderLSTM := NewLSTM(g, embeddingSize, hiddenSize)

	// Decoder embedding
	decoderEmbedding := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(vocabSize, embeddingSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))

	// Decoder LSTM
	decoderLSTM := NewLSTM(g, embeddingSize, hiddenSize)

	// Decoder output dense layer
	decoderDense := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(hiddenSize, vocabSize), gorgonia.WithInit(gorgonia.GlorotN(1.0)))

	return &Seq2Seq{
		encoderEmbedding: encoderEmbedding,
		encoderLSTM:      encoderLSTM,
		decoderEmbedding: decoderEmbedding,
		decoderLSTM:      decoderLSTM,
		decoderDense:     decoderDense,
		vocabSize:        vocabSize,
		embeddingSize:    embeddingSize,
		hiddenSize:       hiddenSize,
	}
}

func (s *Seq2Seq) Learnables() gorgonia.Nodes {
	var nodes gorgonia.Nodes
	nodes = append(nodes, s.encoderEmbedding)
	nodes = append(nodes, s.encoderLSTM.Parameters()...)
	nodes = append(nodes, s.decoderEmbedding)
	nodes = append(nodes, s.decoderLSTM.Parameters()...)
	nodes = append(nodes, s.decoderDense)
	return nodes
}

// Parameters возвращает все обучаемые параметры LSTM слоя
func (lstm *LSTM) Parameters() gorgonia.Nodes {
	return gorgonia.Nodes{
		lstm.wf, lstm.uf, lstm.bf,
		lstm.wi, lstm.ui, lstm.bi,
		lstm.wo, lstm.uo, lstm.bo,
		lstm.wg, lstm.ug, lstm.bg,
	}
}

// Реализация функции эмбеддинга
func embedding(embeddingMatrix *gorgonia.Node, indices *gorgonia.Node, g *gorgonia.ExprGraph) (*gorgonia.Node, error) {
	// Преобразуем индексы в one-hot векторы
	oneHotVectors, err := oneHot(indices, embeddingMatrix.Shape()[0])
	if err != nil {
		return nil, err
	}
	// Вычисляем произведение one-hot векторов и матрицы эмбеддингов
	embedded := gorgonia.Must(gorgonia.Mul(oneHotVectors, embeddingMatrix))
	return embedded, nil
}

// Реализация функции one-hot кодирования
func oneHot(indices *gorgonia.Node, numClasses int) (*gorgonia.Node, error) {
	indicesTensor, ok := indices.Value().(tensor.Tensor)
	if !ok {
		return nil, fmt.Errorf("indices is not a tensor")
	}
	indicesData := indicesTensor.Data()

	var indicesInt []int
	switch data := indicesData.(type) {
	case []int:
		indicesInt = data
	case []int32:
		for _, v := range data {
			indicesInt = append(indicesInt, int(v))
		}
	case []int64:
		for _, v := range data {
			indicesInt = append(indicesInt, int(v))
		}
	default:
		return nil, fmt.Errorf("unsupported indices data type")
	}

	batchSize := len(indicesInt)
	oneHotData := make([]float32, batchSize*numClasses)
	for i, idx := range indicesInt {
		if idx >= numClasses {
			return nil, fmt.Errorf("index %d out of range", idx)
		}
		oneHotData[i*numClasses+idx] = 1.0
	}

	oneHotTensor := tensor.New(tensor.WithShape(batchSize, numClasses), tensor.WithBacking(oneHotData))
	oneHotNode := gorgonia.NewMatrix(indices.Graph(), tensor.Float32, gorgonia.WithShape(batchSize, numClasses), gorgonia.WithValue(oneHotTensor))
	return oneHotNode, nil
}

// Функция для создания one-hot матрицы целей
func oneHotTargets(targets *gorgonia.Node, numClasses int) *gorgonia.Node {
	// targets имеет форму [batchSize]
	targetsTensor := targets.Value().(tensor.Tensor)
	targetsData := targetsTensor.Data().([]int)
	batchSize := len(targetsData)

	oneHotData := make([]float32, batchSize*numClasses)
	for i, idx := range targetsData {
		if idx >= numClasses {
			idx = 0 // Если индекс вне диапазона, заменяем на 0
		}
		oneHotData[i*numClasses+idx] = 1.0
	}

	oneHotTensor := tensor.New(tensor.WithShape(batchSize, numClasses), tensor.WithBacking(oneHotData))
	oneHotNode := gorgonia.NewMatrix(targets.Graph(), tensor.Float32, gorgonia.WithShape(batchSize, numClasses), gorgonia.WithValue(oneHotTensor))
	return oneHotNode
}

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

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

	var encoderInputs [][]int
	var decoderInputs [][]int
	var decoderTargets [][]int

	for _, item := range parsedData {
		inputSeq := Vectorize(Tokenize(item.Input), vocab)
		outputSeq := Vectorize(Tokenize(item.Output), vocab)

		decoderInputSeq := append([]int{vocab[startToken]}, outputSeq...)
		decoderTargetSeq := append(outputSeq, vocab[endToken])

		encoderInputs = append(encoderInputs, inputSeq)
		decoderInputs = append(decoderInputs, decoderInputSeq)
		decoderTargets = append(decoderTargets, decoderTargetSeq)
	}

	// Определяем максимальную длину последовательности
	maxEncoderLen := maxSeqLength(encoderInputs)
	maxDecoderLen := maxSeqLength(decoderInputs)

	// Паддинг последовательностей
	encoderInputsPadded := padSequences(encoderInputs, maxEncoderLen, vocab[padToken])
	decoderInputsPadded := padSequences(decoderInputs, maxDecoderLen, vocab[padToken])
	decoderTargetsPadded := padSequences(decoderTargets, maxDecoderLen, vocab[padToken])

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

	// Создаем граф
	vocabSize := len(vocab)
	model := NewSeq2Seq(nil, vocabSize, embeddingSize, hiddenSize)

	// Оптимизатор
	optimizer := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(0.001))

	totalSamples := len(encoderInputsPadded)
	indices := rand.Perm(totalSamples)

	for epoch := 0; epoch < numEpochs; epoch++ {
		fmt.Printf("Эпоха %d/%d\n", epoch+1, numEpochs)
		var totalLoss float32

		// Перемешиваем данные перед каждой эпохой
		indices = rand.Perm(totalSamples)

		for i := 0; i < totalSamples; i += batchSize {
			end := i + batchSize
			if end > totalSamples {
				end = totalSamples
			}

			batchIndices := indices[i:end]
			encoderBatch := selectBatch(encoderInputsPadded, batchIndices)
			decoderInputBatch := selectBatch(decoderInputsPadded, batchIndices)
			decoderTargetBatch := selectBatch(decoderTargetsPadded, batchIndices)

			// Создаем новый граф для каждого батча
			g := gorgonia.NewGraph()
			model := NewSeq2Seq(g, vocabSize, embeddingSize, hiddenSize)

			// Копируем значения параметров из модели
			shapes := make([]tensor.Shape, len(model.Learnables()))
			for i, param := range model.Learnables() {
				shapes[i] = param.Shape()
			}
			// Загрузка параметров
			params, err := LoadModel(g, shapes, "models/seq2seq_model.gob")
			if err != nil {
				// Если модель еще не сохранена, это не ошибка
				params = model.Learnables()
			}
			for i, param := range params {
				gorgonia.Let(model.Learnables()[i], param.Value())
			}

			loss, err := trainStep(g, model, encoderBatch, decoderInputBatch, decoderTargetBatch, optimizer)
			if err != nil {
				log.Fatalf("Ошибка во время обучения: %v", err)
			}
			totalLoss += loss

			fmt.Printf("Обучение: %.2f%%\r", float32(end)/float32(totalSamples)*100)
		}
		fmt.Printf("Средняя потеря: %.4f\n", totalLoss/float32(totalSamples))

		// Сохраняем модель после каждой эпохи
		if err := SaveModel(model.Learnables(), "models/seq2seq_model.gob"); err != nil {
			log.Fatalf("Ошибка при сохранении модели: %v", err)
		}
	}

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

func trainStep(g *gorgonia.ExprGraph, model *Seq2Seq, encoderInputs, decoderInputs, decoderTargets [][]int, optimizer gorgonia.Solver) (float32, error) {
	batchSize := len(encoderInputs)
	maxEncoderLen := len(encoderInputs[0])
	maxDecoderLen := len(decoderInputs[0])

	// Создаем узлы для входов
	encoderInputNodes := make([]*gorgonia.Node, maxEncoderLen)
	for t := 0; t < maxEncoderLen; t++ {
		encoderInputNodes[t] = gorgonia.NewVector(g, tensor.Int, gorgonia.WithShape(batchSize), gorgonia.WithName(fmt.Sprintf("encoder_input_%d", t)))
	}

	decoderInputNodes := make([]*gorgonia.Node, maxDecoderLen)
	for t := 0; t < maxDecoderLen; t++ {
		decoderInputNodes[t] = gorgonia.NewVector(g, tensor.Int, gorgonia.WithShape(batchSize), gorgonia.WithName(fmt.Sprintf("decoder_input_%d", t)))
	}

	decoderTargetNodes := make([]*gorgonia.Node, maxDecoderLen)
	for t := 0; t < maxDecoderLen; t++ {
		decoderTargetNodes[t] = gorgonia.NewVector(g, tensor.Int, gorgonia.WithShape(batchSize), gorgonia.WithName(fmt.Sprintf("decoder_target_%d", t)))
	}

	// Энкодер
	var encoderHidden *gorgonia.Node
	var encoderCell *gorgonia.Node
	h0 := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
	c0 := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
	encoderHidden = h0
	encoderCell = c0

	for t := 0; t < maxEncoderLen; t++ {
		// Извлекаем входы
		x_t := encoderInputNodes[t]
		// Получаем векторы эмбеддинга
		x_t_embedded, err := embedding(model.encoderEmbedding, x_t, g)
		if err != nil {
			return 0, err
		}

		encoderHidden, encoderCell, _ = model.encoderLSTM.Forward(x_t_embedded, encoderHidden, encoderCell)
	}

	// Декодер
	decoderHidden := encoderHidden
	decoderCell := encoderCell

	// Список для накопления потерь
	losses := make([]*gorgonia.Node, 0, maxDecoderLen)

	for t := 0; t < maxDecoderLen; t++ {
		x_t := decoderInputNodes[t]
		x_t_embedded, err := embedding(model.decoderEmbedding, x_t, g)
		if err != nil {
			return 0, err
		}

		decoderHidden, decoderCell, _ = model.decoderLSTM.Forward(x_t_embedded, decoderHidden, decoderCell)

		logits := gorgonia.Must(gorgonia.Mul(decoderHidden, model.decoderDense)) // [batchSize, vocabSize]
		prediction := gorgonia.Must(gorgonia.SoftMax(logits))

		target := decoderTargetNodes[t]

		// Вычисляем потери с помощью функции негативной логарифмической правдоподобности
		logPredictions := gorgonia.Must(gorgonia.Log(prediction))
		oneHotTargetsNode := oneHotTargets(target, model.vocabSize)
		nllLoss := gorgonia.Must(gorgonia.Neg(gorgonia.Must(gorgonia.Sum(gorgonia.Must(gorgonia.HadamardProd(logPredictions, oneHotTargetsNode)), 1))))
		losses = append(losses, nllLoss)
	}

	// Суммируем потери
	totalLoss := losses[0]
	for i := 1; i < len(losses); i++ {
		totalLoss = gorgonia.Must(gorgonia.Add(totalLoss, losses[i]))
	}
	// Находим среднее значение потерь
	totalLoss = gorgonia.Must(gorgonia.Div(totalLoss, gorgonia.NewConstant(float32(len(losses)))))

	// Создаем градиенты
	if _, err := gorgonia.Grad(totalLoss, model.Learnables()...); err != nil {
		return 0, err
	}

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

	// Задаем значения входов
	for t := 0; t < maxEncoderLen; t++ {
		data := make([]int, batchSize)
		for i := 0; i < batchSize; i++ {
			data[i] = encoderInputs[i][t]
		}
		x_t_val := tensor.New(tensor.WithShape(batchSize), tensor.WithBacking(data))
		if err := gorgonia.Let(encoderInputNodes[t], x_t_val); err != nil {
			return 0, err
		}
	}

	for t := 0; t < maxDecoderLen; t++ {
		// Входы декодера
		data := make([]int, batchSize)
		for i := 0; i < batchSize; i++ {
			data[i] = decoderInputs[i][t]
		}
		x_t_val := tensor.New(tensor.WithShape(batchSize), tensor.WithBacking(data))
		if err := gorgonia.Let(decoderInputNodes[t], x_t_val); err != nil {
			return 0, err
		}

		// Целевые значения
		targets := make([]int, batchSize)
		for i := 0; i < batchSize; i++ {
			targets[i] = decoderTargets[i][t]
		}
		y_t_val := tensor.New(tensor.WithShape(batchSize), tensor.WithBacking(targets))
		if err := gorgonia.Let(decoderTargetNodes[t], y_t_val); err != nil {
			return 0, err
		}
	}

	// Запускаем машину
	if err := machine.RunAll(); err != nil {
		return 0, err
	}

	// Получаем значение потерь
	lossValue := totalLoss.Value().Data().(float32)

	// Обновляем параметры
	if err := optimizer.Step(gorgonia.NodesToValueGrads(model.Learnables())); err != nil {
		return 0, err
	}

	machine.Reset()

	return lossValue, nil
}

func maxSeqLength(sequences [][]int) int {
	maxLen := 0
	for _, seq := range sequences {
		if len(seq) > maxLen {
			maxLen = len(seq)
		}
	}
	return maxLen
}

func padSequences(sequences [][]int, maxLen int, padValue int) [][]int {
	var padded [][]int
	for _, seq := range sequences {
		if len(seq) < maxLen {
			pad := make([]int, maxLen-len(seq))
			for i := range pad {
				pad[i] = padValue
			}
			seq = append(seq, pad...)
		} else if len(seq) > maxLen {
			seq = seq[:maxLen]
		}
		padded = append(padded, seq)
	}
	return padded
}

func selectBatch(sequences [][]int, indices []int) [][]int {
	var batch [][]int
	for _, idx := range indices {
		batch = append(batch, sequences[idx])
	}
	return batch
}

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

	idToWord := make(map[int]string)
	for word, id := range vocab {
		idToWord[id] = word
	}

	vocabSize := len(vocab)
	g := gorgonia.NewGraph()
	model := NewSeq2Seq(g, vocabSize, embeddingSize, hiddenSize)

	// Загрузка обученной модели
	learnables := model.Learnables()
	shapes := make([]tensor.Shape, len(learnables))
	for i, param := range learnables {
		shapes[i] = param.Shape()
	}

	params, err := LoadModel(g, shapes, "models/seq2seq_model.gob")
	if err != nil {
		log.Fatalf("Ошибка при загрузке модели: %v", err)
	}

	// Устанавливаем значения параметров модели
	for i, param := range params {
		gorgonia.Let(learnables[i], param.Value())
	}

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("Введите сообщение (или 'exit' для выхода): ")
		text, _ := reader.ReadString('\n')
		text = strings.TrimSpace(text)

		if text == "exit" {
			break
		}

		inputSeq := Vectorize(Tokenize(text), vocab)
		inputSeq = append(inputSeq, vocab[endToken])

		maxEncoderLen := len(inputSeq)
		maxDecoderLen := 20 // Ограничиваем длину ответа

		// Создаем новый граф для генерации
		g := gorgonia.NewGraph()
		model := NewSeq2Seq(g, vocabSize, embeddingSize, hiddenSize)

		// Устанавливаем параметры модели
		for i, param := range params {
			gorgonia.Let(model.Learnables()[i], param.Value())
		}

		// Энкодер
		var encoderHidden *gorgonia.Node
		var encoderCell *gorgonia.Node
		h0 := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
		c0 := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
		encoderHidden = h0
		encoderCell = c0

		for t := 0; t < maxEncoderLen; t++ {
			// Создаем узел для входа
			x_t := gorgonia.NewScalar(g, tensor.Int)
			if err := gorgonia.Let(x_t, inputSeq[t]); err != nil {
				log.Fatalf("Ошибка при установке значения входа: %v", err)
			}

			// Получаем эмбеддинг
			x_t_embedded, err := embedding(model.encoderEmbedding, x_t, g)
			if err != nil {
				log.Fatalf("Ошибка при вычислении эмбеддинга: %v", err)
			}

			encoderHidden, encoderCell, _ = model.encoderLSTM.Forward(x_t_embedded, encoderHidden, encoderCell)
		}

		// Декодер
		decoderHidden := encoderHidden
		decoderCell := encoderCell

		currentWord := vocab[startToken]
		var generatedSequence []int

		for t := 0; t < maxDecoderLen; t++ {
			// Получаем эмбеддинг текущего слова
			x_t := gorgonia.NewScalar(g, tensor.Int)
			if err := gorgonia.Let(x_t, currentWord); err != nil {
				log.Fatalf("Ошибка при установке текущего слова: %v", err)
			}

			x_t_embedded, err := embedding(model.decoderEmbedding, x_t, g)
			if err != nil {
				log.Fatalf("Ошибка при вычислении эмбеддинга декодера: %v", err)
			}

			decoderHidden, decoderCell, _ = model.decoderLSTM.Forward(x_t_embedded, decoderHidden, decoderCell)

			logits := gorgonia.Must(gorgonia.Mul(decoderHidden, model.decoderDense))
			prediction := gorgonia.Must(gorgonia.SoftMax(logits))

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

			predictedProbs := prediction.Value().Data().([]float32)
			predictedWordID := argmax(predictedProbs)

			if predictedWordID == vocab[endToken] || predictedWordID == vocab[padToken] {
				break
			}

			generatedSequence = append(generatedSequence, predictedWordID)
			currentWord = predictedWordID

			machine.Reset()
		}

		// Преобразуем предсказанные индексы обратно в слова
		var outputTokens []string
		for _, idx := range generatedSequence {
			if word, exists := idToWord[idx]; exists {
				outputTokens = append(outputTokens, word)
			} else {
				outputTokens = append(outputTokens, unkToken)
			}
		}

		response := strings.Join(outputTokens, " ")
		fmt.Printf("Ответ: %s\n", response)
	}
}

func argmax(slice []float32) int {
	maxIdx := 0
	maxVal := slice[0]
	for i, val := range slice {
		if val > maxVal {
			maxVal = val
			maxIdx = i
		}
	}
	return maxIdx
}

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

	switch mode {
	case "t":
		trainModel("fgothi", maxTrainMessages)
	case "d":
		interactWithModel()
	default:
		fmt.Println("Неизвестный режим. Используйте 't' для обучения или 'd' для диалога.")
	}
}

Введите режим работы (t - обучение, d - диалог): t
Загружено 1896 обучающих пар
panic: no graph supplied %!s(<nil>)() :: Matrix float32

goroutine 1 [running]:
gorgonia.org/gorgonia.(*Node).fix(0xc0001be000?)
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/node.go:790 +0xe5
gorgonia.org/gorgonia.newNode({0xc000114540, 0x4, 0x4?})
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/node.go:281 +0xcd
gorgonia.org/gorgonia.NewUniqueNode({0xc000114540?, 0xc000261970?, 0x10?})
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/node.go:289 +0x18
gorgonia.org/gorgonia.NewMatrix(0x0, {{0x164c920, 0x14c8000}}, {0xc000261a20, 0x2, 0xc0002619e8?})
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/gorgonia.go:69 +0x1b2
main.NewSeq2Seq(0x0, 0x544b, 0x80, 0x100)
        D:/JBProjects/GOLAND/Learn/main.go:355 +0x10f
main.trainModel({0x157a85a?, 0xc00010e020?}, 0xc000139f30?)
        D:/JBProjects/GOLAND/Learn/main.go:520 +0xf32
main.main()
        D:/JBProjects/GOLAND/Learn/main.go:916 +0x106