package main

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

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

// ParsedData represents the structure for input and output data
type ParsedData struct {
	Input  string
	Output string
}

// Constants
const (
	maxTrainMessages = 10000
	maxVocabSize     = 50000
	numEpochs        = 10
	batchSize        = 64
	embeddingSize    = 128
	hiddenSize       = 256
	startToken       = "<SOS>"
	endToken         = "<EOS>"
	padToken         = "<PAD>"
	unkToken         = "<UNK>"
)

// LoadFiles loads files from the specified directory
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
}

// SaveVocabulary saves the vocabulary to a file
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 loads the vocabulary from a file
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
}

// SaveModel saves the model parameters to a file
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)
}

// LoadModel loads the model parameters from a file
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 model
type Seq2Seq struct {
	encoderLSTM *LSTM

	decoderLSTM  *LSTM
	decoderDense *gorgonia.Node

	vocabSize     int
	embeddingSize int
	hiddenSize    int

	// Precomputed embedding matrices (not trainable within the graph)
	encoderEmbeddings [][]float32
	decoderEmbeddings [][]float32
}

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

// NewLSTM creates a new LSTM layer
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,
	}
}

// Parameters returns all trainable parameters of the LSTM layer
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,
	}
}

// Forward performs forward pass through the LSTM layer
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
}

// NewSeq2Seq creates a new Seq2Seq model
func NewSeq2Seq(g *gorgonia.ExprGraph, vocabSize, embeddingSize, hiddenSize int, encoderEmbeddings, decoderEmbeddings [][]float32) *Seq2Seq {
	// Encoder LSTM
	encoderLSTM := NewLSTM(g, embeddingSize, hiddenSize)

	// 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{
		encoderLSTM:       encoderLSTM,
		decoderLSTM:       decoderLSTM,
		decoderDense:      decoderDense,
		vocabSize:         vocabSize,
		embeddingSize:     embeddingSize,
		hiddenSize:        hiddenSize,
		encoderEmbeddings: encoderEmbeddings,
		decoderEmbeddings: decoderEmbeddings,
	}
}

// Learnables returns all trainable parameters of the model
func (s *Seq2Seq) Learnables() gorgonia.Nodes {
	var nodes gorgonia.Nodes
	nodes = append(nodes, s.encoderLSTM.Parameters()...)
	nodes = append(nodes, s.decoderLSTM.Parameters()...)
	nodes = append(nodes, s.decoderDense)
	return nodes
}

// GenerateRandomEmbeddings generates random embeddings for the vocabulary
func GenerateRandomEmbeddings(vocabSize, embeddingSize int) [][]float32 {
	embeddings := make([][]float32, vocabSize)
	for i := 0; i < vocabSize; i++ {
		vec := make([]float32, embeddingSize)
		for j := 0; j < embeddingSize; j++ {
			vec[j] = rand.Float32()*0.2 - 0.1 // Random values between -0.1 and 0.1
		}
		embeddings[i] = vec
	}
	return embeddings
}

func trainModel(trainDir string, maxTrainMessages int) {
	parsedData, err := LoadFiles(trainDir, maxTrainMessages)
	if err != nil {
		log.Fatalf("Error loading training files: %v", err)
	}

	fmt.Printf("Loaded %d training pairs\n", len(parsedData))

	vocab := BuildVocab(parsedData)

	// Create models directory if it doesn't exist
	if err := os.MkdirAll("models", os.ModePerm); err != nil {
		log.Fatalf("Error creating models directory: %v", err)
	}

	// Save vocabulary to models directory
	if err := SaveVocabulary(vocab, "models/vocab.gob"); err != nil {
		log.Fatalf("Error saving vocabulary: %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)
	}

	// Determine maximum sequence length
	maxEncoderLen := maxSeqLength(encoderInputs)
	maxDecoderLen := maxSeqLength(decoderInputs)

	// Pad sequences
	encoderInputsPadded := padSequences(encoderInputs, maxEncoderLen, vocab[padToken])
	decoderInputsPadded := padSequences(decoderInputs, maxDecoderLen, vocab[padToken])
	decoderTargetsPadded := padSequences(decoderTargets, maxDecoderLen, vocab[padToken])

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

	// Create random embeddings (not trainable within the graph)
	vocabSize := len(vocab)
	encoderEmbeddings := GenerateRandomEmbeddings(vocabSize, embeddingSize)
	decoderEmbeddings := GenerateRandomEmbeddings(vocabSize, embeddingSize)

	// Create graph
	g := gorgonia.NewGraph()

	model := NewSeq2Seq(g, vocabSize, embeddingSize, hiddenSize, encoderEmbeddings, decoderEmbeddings)

	// Initialize model parameters
	learnables := model.Learnables()
	shapes := make([]tensor.Shape, len(learnables))
	for i, param := range learnables {
		shapes[i] = param.Shape()
	}

	// Try to load previously saved parameters
	params, err := LoadModel(g, shapes, "models/seq2seq_model.gob")
	if err == nil {
		// If loaded successfully, set model parameters
		for i, param := range params {
			gorgonia.Let(learnables[i], param.Value())
		}
	}

	// Optimizer
	optimizer := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(0.001))

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

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

		// Shuffle data before each epoch
		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)

			loss, err := trainStep(model, encoderBatch, decoderInputBatch, decoderTargetBatch, optimizer)
			if err != nil {
				log.Fatalf("Error during training: %v", err)
			}
			totalLoss += loss

			fmt.Printf("Training: %.2f%%\r", float32(end)/float32(totalSamples)*100)
		}
		fmt.Printf("Average Loss: %.4f\n", totalLoss/float32(totalSamples))

		// Save model after each epoch
		if err := SaveModel(model.Learnables(), "models/seq2seq_model.gob"); err != nil {
			log.Fatalf("Error saving model: %v", err)
		}
	}

	fmt.Println("Training complete.")
}

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

	// Encoder inputs (embedded)
	encoderInputData := make([]float32, batchSize*maxEncoderLen*model.embeddingSize)
	for i := 0; i < batchSize; i++ {
		for t := 0; t < maxEncoderLen; t++ {
			wordIdx := encoderInputs[i][t]
			embedding := model.encoderEmbeddings[wordIdx]
			copy(encoderInputData[(i*maxEncoderLen+t)*model.embeddingSize:(i*maxEncoderLen+t+1)*model.embeddingSize], embedding)
		}
	}
	encoderInputsTensor := tensor.New(tensor.WithShape(batchSize, maxEncoderLen, model.embeddingSize), tensor.WithBacking(encoderInputData))
	xEnc := gorgonia.NewTensor(g, tensor.Float32, 3, gorgonia.WithShape(batchSize, maxEncoderLen, model.embeddingSize), gorgonia.WithValue(encoderInputsTensor))

	// Decoder inputs (embedded)
	decoderInputData := make([]float32, batchSize*maxDecoderLen*model.embeddingSize)
	for i := 0; i < batchSize; i++ {
		for t := 0; t < maxDecoderLen; t++ {
			wordIdx := decoderInputs[i][t]
			embedding := model.decoderEmbeddings[wordIdx]
			copy(decoderInputData[(i*maxDecoderLen+t)*model.embeddingSize:(i*maxDecoderLen+t+1)*model.embeddingSize], embedding)
		}
	}
	decoderInputsTensor := tensor.New(tensor.WithShape(batchSize, maxDecoderLen, model.embeddingSize), tensor.WithBacking(decoderInputData))
	xDec := gorgonia.NewTensor(g, tensor.Float32, 3, gorgonia.WithShape(batchSize, maxDecoderLen, model.embeddingSize), gorgonia.WithValue(decoderInputsTensor))

	// Decoder targets
	decoderTargetsData := make([]int, batchSize*maxDecoderLen)
	for i := 0; i < batchSize; i++ {
		copy(decoderTargetsData[i*maxDecoderLen:(i+1)*maxDecoderLen], decoderTargets[i])
	}
	decoderTargetsTensor := tensor.New(tensor.WithShape(batchSize, maxDecoderLen), tensor.WithBacking(decoderTargetsData))
	yTrue := gorgonia.NewTensor(g, tensor.Int, 2, gorgonia.WithShape(batchSize, maxDecoderLen), gorgonia.WithValue(decoderTargetsTensor))

	// Encoder
	var encoderHidden *gorgonia.Node
	var encoderCell *gorgonia.Node
	h0 := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, model.hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
	c0 := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, model.hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
	encoderHidden = h0
	encoderCell = c0

	// Unstack xEnc along time axis
	xEncSteps := make([]*gorgonia.Node, maxEncoderLen)
	for t := 0; t < maxEncoderLen; t++ {
		x_t := gorgonia.Must(gorgonia.Slice(xEnc, gorgonia.S(0, batchSize), gorgonia.S(t)))
		xEncSteps[t] = x_t
	}

	for t := 0; t < maxEncoderLen; t++ {
		x_t := xEncSteps[t]
		encoderHidden, encoderCell, _ = model.encoderLSTM.Forward(x_t, encoderHidden, encoderCell)
	}

	// Decoder initial hidden state
	decoderHidden := encoderHidden
	decoderCell := encoderCell

	// Unstack xDec and yTrue along time axis
	xDecSteps := make([]*gorgonia.Node, maxDecoderLen)
	yTrueSteps := make([]*gorgonia.Node, maxDecoderLen)
	for t := 0; t < maxDecoderLen; t++ {
		x_t := gorgonia.Must(gorgonia.Slice(xDec, gorgonia.S(0, batchSize), gorgonia.S(t)))
		xDecSteps[t] = x_t

		y_t := gorgonia.Must(gorgonia.Slice(yTrue, gorgonia.S(0, batchSize), gorgonia.S(t)))
		yTrueSteps[t] = y_t
	}

	// Calculate losses for each time step
	var losses []*gorgonia.Node
	for t := 0; t < maxDecoderLen; t++ {
		x_t := xDecSteps[t]
		y_t := yTrueSteps[t]

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

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

		// Convert y_t (shape [batchSize]) to int slice
		y_t_tensor := y_t.Value().(tensor.Tensor)
		y_t_data := y_t_tensor.Data().([]int)

		// Compute NLL loss manually
		predictionTensor := prediction.Value().(tensor.Tensor)
		predictionData := predictionTensor.Data().([]float32)
		vocabSize := model.vocabSize

		nllLossData := make([]float32, batchSize)
		for i := 0; i < batchSize; i++ {
			targetIdx := y_t_data[i]
			predIdx := i * vocabSize
			predProb := predictionData[predIdx+targetIdx]

			// Avoid log(0)
			if predProb < 1e-7 {
				predProb = 1e-7
			}

			nllLossData[i] = -float32(math.Log(float64(predProb)))
		}

		nllLossTensor := tensor.New(tensor.WithShape(batchSize), tensor.WithBacking(nllLossData))
		nllLossNode := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(batchSize), gorgonia.WithValue(nllLossTensor))
		losses = append(losses, nllLossNode)
	}

	// Total loss
	totalLoss := losses[0]
	for i := 1; i < len(losses); i++ {
		totalLoss = gorgonia.Must(gorgonia.Add(totalLoss, losses[i]))
	}
	totalLoss = gorgonia.Must(gorgonia.Mean(totalLoss))

	// Compute gradients
	if _, err := gorgonia.Grad(totalLoss, model.Learnables()...); err != nil {
		return 0, err
	}

	// Create VM
	machine := gorgonia.NewTapeMachine(g)
	defer machine.Close()

	// Run machine
	if err := machine.RunAll(); err != nil {
		return 0, err
	}

	// Get loss value
	lossValue := totalLoss.Value().Data().(float32)

	// Update parameters
	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("Error loading vocabulary: %v", err)
	}

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

	// Create random embeddings (same ones used during training)
	vocabSize := len(vocab)
	encoderEmbeddings := GenerateRandomEmbeddings(vocabSize, embeddingSize)
	decoderEmbeddings := GenerateRandomEmbeddings(vocabSize, embeddingSize)

	g := gorgonia.NewGraph()
	model := NewSeq2Seq(g, vocabSize, embeddingSize, hiddenSize, encoderEmbeddings, decoderEmbeddings)

	// Load trained model parameters
	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("Error loading model: %v", err)
	}

	// Set model parameters
	for i, param := range params {
		gorgonia.Let(learnables[i], param.Value())
	}

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("Enter a message (or 'exit' to quit): ")
		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 // Limit the reply length

		// Prepare encoder inputs
		encoderInputData := make([]float32, maxEncoderLen*model.embeddingSize)
		for t := 0; t < maxEncoderLen; t++ {
			wordIdx := inputSeq[t]
			embedding := model.encoderEmbeddings[wordIdx]
			copy(encoderInputData[t*model.embeddingSize:(t+1)*model.embeddingSize], embedding)
		}
		encoderInputsTensor := tensor.New(tensor.WithShape(1, maxEncoderLen, model.embeddingSize), tensor.WithBacking(encoderInputData))
		xEnc := gorgonia.NewTensor(g, tensor.Float32, 3, gorgonia.WithShape(1, maxEncoderLen, model.embeddingSize), gorgonia.WithValue(encoderInputsTensor))

		// Encoder
		var encoderHidden *gorgonia.Node
		var encoderCell *gorgonia.Node
		h0 := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(model.hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
		c0 := gorgonia.NewVector(g, tensor.Float32, gorgonia.WithShape(model.hiddenSize), gorgonia.WithInit(gorgonia.Zeroes()))
		encoderHidden = h0
		encoderCell = c0

		// Unstack xEnc along time axis
		xEncSteps := make([]*gorgonia.Node, maxEncoderLen)
		for t := 0; t < maxEncoderLen; t++ {
			x_t := gorgonia.Must(gorgonia.Slice(xEnc, gorgonia.S(0), gorgonia.S(t)))
			xEncSteps[t] = x_t
		}

		for t := 0; t < maxEncoderLen; t++ {
			x_t := xEncSteps[t]
			encoderHidden, encoderCell, _ = model.encoderLSTM.Forward(x_t, encoderHidden, encoderCell)
		}

		// Decoder
		decoderHidden := encoderHidden
		decoderCell := encoderCell

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

		for t := 0; t < maxDecoderLen; t++ {
			// Get embedding of current word
			embedding := model.decoderEmbeddings[currentWord]
			x_t_val := tensor.New(tensor.WithShape(1, model.embeddingSize), tensor.WithBacking(embedding))
			x_t := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, model.embeddingSize), gorgonia.WithValue(x_t_val))

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

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

			// Create VM
			machine := gorgonia.NewTapeMachine(g)
			if err := machine.RunAll(); err != nil {
				log.Fatalf("Error during generation: %v", err)
			}

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

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

			generatedSequence = append(generatedSequence, predictedWordID)
			currentWord = predictedWordID

			machine.Reset()
		}

		// Convert predicted indices back to words
		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("Response: %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("Enter mode (t - training, d - dialogue): ")
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	mode := scanner.Text()

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

Enter mode (t - training, d - dialogue): t
Loaded 1896 training pairs
Epoch 1/10
panic: Not all children have the same graph

goroutine 1 [running]:
gorgonia.org/gorgonia.Must(...)
        C:/Users/79685/go/pkg/mod/gorgonia.org/gorgonia@v0.9.18/gorgonia.go:19
main.(*LSTM).Forward(0xc0001c4000, 0xc000195500, 0xc000195340, 0xc000195340)
        D:/JBProjects/GOLAND/Learn/main.go:319 +0x5d6
main.trainStep(0xc00016c0c0, {0xc0000a7c08, 0x40, 0x2?}, {0xc0000a8308, 0x40, 0x544b?}, {0xc0000a8a08, 0x40, 0x4a}, ...)
        D:/JBProjects/GOLAND/Learn/main.go:584 +0xb98
main.trainModel({0x81a89a?, 0xc000072028?}, 0xc0000a3f30?)
        D:/JBProjects/GOLAND/Learn/main.go:510 +0x1bd7
main.main()
        D:/JBProjects/GOLAND/Learn/main.go:865 +0x106