package main

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

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

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

// 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) < 7 {
					continue
				}
				var responseID float32
				fmt.Sscanf(record[5], "%f", &responseID)
				data = append(data, ParsedData{
					Input:  record[4],
					Output: responseID,
				})
			}
		}
	}
	return data, nil
}

// Tokenize простая токенизация текста
func Tokenize(text string) []string {
	return strings.Fields(text)
}

// Vectorize превращает текст в числовой формат (очень простая векторизация)
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
}

// Сохранение весов
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
}

// Загрузка весов
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 trainModel(trainDir string) {
	parsedData, err := LoadCSVFiles(trainDir)
	if err != nil {
		log.Fatalf("Ошибка загрузки CSV файлов: %v", err)
	}

	fmt.Printf("Загружено %d обучающих записей\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)
			}
		}
	}

	batchSize := 100
	totalSamples := len(parsedData)

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

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

	// Use Adam optimizer
	optimizer := gorgonia.NewAdamSolver(gorgonia.WithLearnRate(0.15))

	// Process data in batches
	for epoch := 0; epoch < 10; epoch++ { // 10 epochs
		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("Не удалось создать вычислительный граф: %v", err)
			}

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

			if err := machine.RunAll(); err != nil {
				log.Fatalf("Ошибка запуска машины: %v", err)
			}

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

			machine.Reset() // Reset after each batch

			progress := float64(end) / float64(totalSamples) * 100
			fmt.Printf("Epoch %d: Обработан пакет с %d по %d (%.2f%% обучения завершено)\n", epoch+1, start, end, progress)
		}
	}

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

	if err := SaveWeights(w.Value().(*tensor.Dense), "model_weights.gob"); err != nil {
		log.Fatalf("Ошибка сохранения весов модели: %v", err)
	} else {
		fmt.Println("Весы модели успешно сохранены.")
	}
}

func validateModel(validDir string) {
	parsedData, err := LoadCSVFiles(validDir)
	if err != nil {
		log.Fatalf("Ошибка загрузки CSV файлов: %v", err)
	}

	fmt.Printf("Загружено %d проверочных записей\n", len(parsedData))

	weights, err := LoadWeights("model_weights.gob")
	if err != nil {
		log.Fatalf("Ошибка загрузки весов модели: %v", err)
	}

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

	// Create Graph for validation
	g := gorgonia.NewGraph()

	// Evaluate the model
	batchSize := 100
	totalSamples := len(parsedData)

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

	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)
		if err := machine.RunAll(); err != nil {
			log.Fatalf("Ошибка запуска машины: %v", err)
		}

		// No need for error check here, direct access to data
		predictionsVal := yhat.Value().Data().([]float32)

		// Calculate accuracy
		correct := 0
		for i := 0; i < end-start; i++ {
			if int(predictionsVal[i]+0.5) == int(yReal[i]) {
				correct++
			}
		}
		accuracy := float64(correct) / float64(end-start) * 100
		fmt.Printf("Пакет с %d по %d: Точность = %.2f%%\n", start, end, accuracy)

		machine.Reset()
	}
}

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

	switch mode {
	case "t":
		trainModel("train")
	case "v":
		validateModel("valid")
	default:
		fmt.Println("Неизвестный режим работы. Используйте 't' для тренировки или 'v' для валидации.")
	}
}


сделай так чтобы использовалась видеокарта