document.addEventListener('DOMContentLoaded', function() {
    // Check for dataset parameter in URL
    const urlParams = new URLSearchParams(window.location.search);
    const datasetType = urlParams.get('dataset');
    
    if (datasetType) {
        const datasets = {
            common: "hello world\nhow are you\nwhat is your name\nthis is a test\ngood morning\ngood night\nthank you\nplease wait\nI love coding\nthe quick brown fox\njumps over the lazy dog",
            chat: "brb\nlol\nomg\nttyl\nbtw\nidk\nrofl\nasap\nthx\ncya\nnp\ngtg\nhbu\nimo\nfyi",
            tech: "javascript\npython\nreact\nnodejs\napi\ndatabase\nfunction\nvariable\narray\nobject\nloop\nconditional\nasync\nawait",
            names: "john\nmary\ndavid\nsarah\nmichael\nemma\njames\nolivia\nrobert\nsophia\nnew york\nlondon\nparis\ntokyo\nberlin"
        };
        
        document.getElementById('dataset').value = datasets[datasetType] || '';
    }
const trainBtn = document.getElementById('trainBtn');
    const predictBtn = document.getElementById('predictBtn');
    const datasetTextarea = document.getElementById('dataset');
    const inputWord = document.getElementById('inputWord');
    const resultContainer = document.getElementById('resultContainer');
    const resultText = document.getElementById('resultText');
    const vocabularyList = document.getElementById('vocabularyList');
    const wordCount = document.getElementById('wordCount');
    const modelControls = document.getElementById('modelControls');
    
    // Neural Network model parameters
    const config = {
        hiddenSize: 16,
        learningRate: 0.01,
        iterations: 100
    };
    
    let vocabulary = new Set();
    let model;
    let encoder;
    let isTraining = false;
    
    // Initialize the model when Train button is clicked
    trainBtn.addEventListener('click', async function() {
        if (isTraining) return;
        
        const sentences = datasetTextarea.value
            .split('\n')
            .filter(line => line.trim() !== '');
        
        if (sentences.length === 0) {
            alert('Please enter some training data first!');
            return;
        }
        
        isTraining = true;
        trainBtn.disabled = true;
        trainBtn.innerHTML = '<div class="loading-spinner"></div> Training...';
        
        try {
            // Extract words from sentences
            vocabulary = extractVocabulary(sentences);
            updateVocabularyDisplay();
            
            // Create encoder (word to vector)
            encoder = createEncoder(vocabulary);
            
            // Train the model
            model = await trainModel(sentences, vocabulary, encoder, config);
            
            // Show the prediction controls
            modelControls.classList.remove('hidden');
            resultContainer.classList.add('hidden');
            
            // Show success message
            const originalText = trainBtn.textContent;
            trainBtn.innerHTML = '<i data-feather="check-circle" class="mr-2"></i> Model Trained!';
            setTimeout(() => {
                trainBtn.innerHTML = '<i data-feather="cpu" class="mr-2"></i> Train Model';
                feather.replace();
            }, 2000);
        } catch (error) {
            console.error('Training error:', error);
            alert('Error during training: ' + error.message);
        } finally {
            isTraining = false;
            trainBtn.disabled = false;
            feather.replace();
        }
    });
    
    // Make prediction when Predict button is clicked
    predictBtn.addEventListener('click', function() {
        if (!model) {
            alert('Please train the model first!');
            return;
        }
        
        const typoWord = inputWord.value.trim().toLowerCase();
        if (typoWord === '') {
            alert('Please enter a word to predict');
            return;
        }
        
        try {
            // Predict the most likely correct word
            const prediction = predictWord(typoWord, vocabulary, encoder, model);
            
            // Display the result
            resultText.textContent = `The correct word for "${typoWord}" might be: "${prediction}"`;
            resultContainer.classList.remove('hidden');
            resultContainer.classList.add('fade-in');
        } catch (error) {
            console.error('Prediction error:', error);
            resultText.textContent = `Error: ${error.message}`;
            resultContainer.classList.remove('hidden');
        }
    });
    
    // Helper function to extract vocabulary from sentences
    function extractVocabulary(sentences) {
        const words = new Set();
        sentences.forEach(sentence => {
            sentence.split(/\s+/).forEach(word => {
                const cleanWord = word.toLowerCase().replace(/[^a-z]/g, '');
                if (cleanWord.length > 0) {
                    words.add(cleanWord);
                }
            });
        });
        return words;
    }
    
    // Update the vocabulary display in the UI
    function updateVocabularyDisplay() {
        vocabularyList.innerHTML = '';
        Array.from(vocabulary).sort().forEach(word => {
            const wordEl = document.createElement('div');
            wordEl.className = 'word-badge px-3 py-1 bg-blue-100 text-blue-800 rounded-full text-sm';
            wordEl.textContent = word;
            vocabularyList.appendChild(wordEl);
        });
        wordCount.textContent = vocabulary.size + ' words';
    }
    
    // Create encoder (simple character-based encoding)
    function createEncoder(vocabulary) {
        const allWords = Array.from(vocabulary);
        const allChars = new Set();
        
        allWords.forEach(word => {
            word.split('').forEach(char => allChars.add(char));
        });
        
        const charToIndex = {};
        Array.from(allChars).sort().forEach((char, index) => {
            charToIndex[char] = index;
        });
        
        return {
            encode: function(word) {
                // Simple bag-of-chars encoding
                const encoded = new Array(allChars.size).fill(0);
                word.split('').forEach(char => {
                    if (charToIndex[char] !== undefined) {
                        encoded[charToIndex[char]] += 1;
                    }
                });
                return encoded;
            },
            maxLength: Math.max(...allWords.map(w => w.length))
        };
    }
    
    // Train the model
    function trainModel(sentences, vocabulary, encoder, config) {
        return new Promise((resolve) => {
            // Simple neural network (simulated)
            setTimeout(() => {
                resolve({
                    predict: function(input) {
                        // Simulate prediction by finding the closest word in vocabulary
                        const inputEncoding = encoder.encode(input);
                        let minDistance = Infinity;
                        let bestMatch = input;
                        
                        vocabulary.forEach(word => {
                            const wordEncoding = encoder.encode(word);
                            const distance = calculateDistance(inputEncoding, wordEncoding);
                            if (distance < minDistance) {
                                minDistance = distance;
                                bestMatch = word;
                            }
                        });
                        
                        return bestMatch;
                    }
                });
            }, 1000); // Simulate training time
        });
    }
    
    // Helper function to calculate distance between encodings
    function calculateDistance(a, b) {
        let distance = 0;
        for (let i = 0; i < a.length; i++) {
            distance += Math.abs(a[i] - b[i]);
        }
        return distance;
    }
    
    // Predict the most likely correct word
    function predictWord(input, vocabulary, encoder, model) {
        return model.predict(input);
    }
});