package main

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

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

	"C"
)

// 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 := 150
	totalSamples := len(parsedData)

	// Use CUDA for Graph, если поддержка CUDA реализована корректно
	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(99.9))

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

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

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

	batchSize := 100
	totalSamples := len(parsedData)

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

	// Validate in batches
	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("Ошибка запуска машины: %v", err)
		}

		predictionsVal := yhat.Value().Data().([]float32)

		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' для валидации.")
	}
}



добавь переменнуб maxTrainmes которое обозночает макисмальное количество сообзений для обучения если эьо число больше всех сообщений то тогда беруться все найденные сообещния если меньше то тогда береться ровно столько сообщений сколько ты указал в переменной