package main

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

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

// ParsedData represents a question and answer pair
type ParsedData struct {
	Input  string
	Output float32
}

const maxTrainMessages = 50000

// LoadCSVFiles loads and processes CSV files, extracting question-answer pairs.
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++ { // skip the header
				record := records[i]
				if len(record) < 7 {
					continue
				}
				var responseID float32
				fmt.Sscanf(record[5], "%f", &responseID)
				data = append(data, ParsedData{
					Input:  record[4],
					Output: responseID,
				})
			}
		}
	}
	return data, nil
}

// Tokenize is an improved text tokenization function
func Tokenize(text string) []string {
	// Remove punctuation and convert to lowercase
	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))
}

// Vectorize converts text to a numeric format (very simple vectorization)
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]++
		}
	}
	return vector
}

// SaveWeights saves the model weights
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)
	if err := encoder.Encode(weights); err != nil {
		return err
	}
	return nil
}

// LoadWeights loads the model 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
}

// 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)
	if err := encoder.Encode(vocab); err != nil {
		return err
	}
	return nil
}

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

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 := make(map[string]int)
	for _, item := range parsedData {
		tokens := Tokenize(item.Input)
		for _, token := range tokens {
			if _, exists := vocab[token]; !exists {
				vocab[token] = len(vocab)
			}
		}
	}

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

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

			xVal := tensor.New(tensor.WithShape(batchSize, inputSize), tensor.Of(tensor.Float32))
			yVal := tensor.New(tensor.WithShape(batchSize, 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)
				yVal.SetAt(parsedData[i].Output, i-start, 0)
			}

			x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, inputSize), gorgonia.WithValue(xVal))
			y := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, 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()

			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(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 validateModel(validDir string) {
	parsedData, err := LoadCSVFiles(validDir)
	if err != nil {
		log.Fatalf("Error loading CSV files: %v", err)
	}

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

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

	g := gorgonia.NewGraph()

	batchSize := 100
	totalSamples := len(parsedData)

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

	correct := 0
	total := 0

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

		xVal := tensor.New(tensor.WithShape(batchSize, inputSize), tensor.Of(tensor.Float32))
		yReal := make([]float32, batchSize)

		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)
			yReal[i-start] = parsedData[i].Output
		}

		x := gorgonia.NewMatrix(g, tensor.Float32, gorgonia.WithShape(batchSize, 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)

		for i := 0; i < end-start; i++ {
			if int(predictionsVal[i]+0.5) == int(yReal[i]) {
				correct++
			}
			total++
		}

		machine.Reset()
	}

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

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

	// Reverse map to get words from vocabulary indices
	reverseVocab := make(map[int]string)
	for word, idx := range vocab {
		reverseVocab[idx] = word
	}

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

	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)

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

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

		// Get the scalar result from the prediction tensor
		predictionsVal := yhat.Value().Data().([]float32)

		topIndex := getTopIndex(predictionsVal) // Get the index of the highest value

		response := "I am not sure how to answer that."

		if word, found := reverseVocab[topIndex]; found {
			response = word
		}

		fmt.Printf("Generated Response: %s\n", response)

		machine.Reset()
	}
}

// getTopIndex returns the index of the highest value in a float slice
func getTopIndex(values []float32) int {
	maxIdx := 0
	maxVal := float32(-1) // Assuming output will not be negative

	for i, val := range values {
		if val > maxVal {
			maxVal = val
			maxIdx = i
		}
	}

	return maxIdx
}

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


Ask a question (or type 'exit' to quit): hello
Generated Response: 115712
Ask a question (or type 'exit' to quit): 

исправь ИИ доолжен выдавать тектовый запрос а не цифры он доден сам составлять запрос перепиши все нужные фунции может нужно переписать валидацию или траининг сам реши но чтобы рабоатло