package main

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

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

type ParsedData struct {
	Input  string
	Output string
}

const maxTrainMessages = 10000
const maxVocabSize = 500000

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
}

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

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
}

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 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 := BuildVocab(parsedData)

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

	batchSize := 150
	totalSamples := len(parsedData)

	g := gorgonia.NewGraph()

	inputSize := len(vocab)
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, 1), gorgonia.WithInit(gorgonia.GlorotN(1.0)))

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

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

	totalProcessed := 0

	for epoch := 0; epoch < 10; epoch++ {
		for start := 0; start < totalSamples; start += batchSize {
			end := start + batchSize
			if end > totalSamples {
				end = totalSamples
			}

			batchSizeActual := end - start
			xVal := tensor.New(tensor.WithShape(batchSizeActual, inputSize), tensor.Of(tensor.Float32))
			yVal := tensor.New(tensor.WithShape(batchSizeActual, 1), tensor.Of(tensor.Float32))

			for i := start; i < end; i++ {
				tokens := Tokenize(parsedData[i].Input)
				inputVec := Vectorize(tokens, vocab)
				offset := (i - start) * inputSize
				copy(xVal.Data().([]float32)[offset:], inputVec)

				// Directly use the output as a float32 value
				outputValue := float32(0) // Default to 0. Set an appropriate default if your data has a specific format.
				if parsedData[i].Output != "" {
					outputValue = float32(parsedData[i].Output[0]) // Adjust this according to how you want to handle the Output
				}
				yVal.SetAt(outputValue, i-start, 0)
			}

			x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSizeActual, inputSize), gorgonia.WithValue(xVal))
			y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSizeActual, 1), gorgonia.WithValue(yVal))

			yhat := gorgonia.Must(gorgonia.Mul(x, w))
			loss := gorgonia.Must(gorgonia.Mean(gorgonia.Must(gorgonia.Square(gorgonia.Must(gorgonia.Sub(yhat, y))))))

			if _, err := gorgonia.Grad(loss, w); err != nil {
				log.Fatalf("Failed to create computation graph: %v", err)
			}

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

			optimizer.Step(gorgonia.NodesToValueGrads([]*gorgonia.Node{w}))

			machine.Reset()

			totalProcessed += batchSizeActual
			progress := float64(start+batchSizeActual) / float64(totalSamples) * 100
			cumulativeProgress := float64(totalProcessed) / float64(totalSamples*10) * 100

			fmt.Printf("Epoch %d: Processed batch from %d to %d (%.2f%% epoch progress, %.2f%% overall progress)\n", epoch+1, start, end, progress, cumulativeProgress)
		}
	}

	fmt.Println("Model training complete")

	if err := SaveWeights(w.Value().(*tensor.Dense), "model_weights.gob"); err != nil {
		log.Fatalf("Error saving model weights: %v", err)
	} else {
		fmt.Println("Model weights saved successfully.")
	}
}

func interactWithModel() {
	weightsTensor, 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)
	}

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

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

	g := gorgonia.NewGraph()
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(inputSize, 1), gorgonia.WithValue(weightsTensor))

	for {
		fmt.Print("Ask a question (or type 'exit' to quit): ")
		text, _ := reader.ReadString('\n')
		text = strings.TrimSpace(text)

		if text == "exit" {
			break
		}

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

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

		x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, inputSize), gorgonia.WithValue(xVal))
		yhat := gorgonia.Must(gorgonia.Mul(x, w))

		machine := gorgonia.NewTapeMachine(g)
		defer machine.Close()

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

		predictionsVal := yhat.Value().Data().([]float32)
		topIndices := getTopIndices(predictionsVal, 5)

		var responseTokens []string
		for _, idx := range topIndices {
			if word, exists := reverseVocab[idx]; exists {
				responseTokens = append(responseTokens, word)
			}
		}

		response := strings.Join(responseTokens, " ")
		fmt.Printf("Generated Response: %s\n", response)
		machine.Reset()
	}
}

func validateModel(validationDir string) {
	weightsTensor, 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)
	}

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

	g := gorgonia.NewGraph()
	w := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(len(vocab), 1), gorgonia.WithValue(weightsTensor))

	correct := 0
	total := len(validationData)

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

		xVal := tensor.New(tensor.WithShape(1, len(vocab)), tensor.Of(tensor.Float32))
		copy(xVal.Data().([]float32), inputVec)

		x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(1, len(vocab)), gorgonia.WithValue(xVal))
		yhat := gorgonia.Must(gorgonia.Mul(x, w))

		machine := gorgonia.NewTapeMachine(g)
		defer machine.Close()

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

		predictionsVal := yhat.Value().Data().([]float32)
		outputResult := 0
		if strings.TrimSpace(item.Output) == "1" {
			outputResult = 1
		}

		if int(predictionsVal[0]+0.5) == outputResult {
			correct++
		}

		machine.Reset()
	}

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

func getTopIndices(values []float32, n int) []int {
	type kv struct {
		Key   int
		Value float32
	}

	var ss []kv
	for i, val := range values {
		ss = append(ss, kv{i, val})
	}

	sort.Slice(ss, func(i, j int) bool {
		return ss[i].Value > ss[j].Value
	})

	var indices []int
	for i := 0; i < n && i < len(ss); i++ {
		indices = append(indices, ss[i].Key)
	}
	return indices
}

func main() {
	fmt.Print("Enter mode of operation (t for training, v for validation, d for dialogue): ")
	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("Unknown mode. Use 't' for training, 'v' for validation, or 'd' for dialogue.")
	}
}


исправь валидацию модели ведь в аппке valid может быть как просто csv файлы так и txt файл и много разных папок в которых тоже етсь txt или csv файлы

исправь ведь валидация всегда на 100%

и измени ткенизацию как для валидации так и для обучения trsin чтобы было так [сообщение человека]: [ответ другого человека на сообщение первого]