"""Training for Veda Programming Assistant"""

import tensorflow as tf
from tensorflow import keras
import numpy as np
import os
import json

from model import VedaProgrammingLLM
from tokenizer import VedaTokenizer
from config import MODEL_DIR, VOCAB_SIZE, MAX_LENGTH, BATCH_SIZE


TRAINING_DATA = '''
<USER> Hello!
<ASSISTANT> Hello! Welcome to Veda Programming Assistant! I can help you with coding, explain concepts, or answer programming questions. What would you like to learn today?

<USER> What is Python?
<ASSISTANT> Python is a high-level, interpreted programming language created by Guido van Rossum in 1991. It is known for its simple, readable syntax.

Key features:
- Easy to learn and read
- Dynamic typing
- Large standard library
- Great for beginners and experts

Python is used for web development, data science, AI, automation, and more!

<USER> How do I print hello world?
<ASSISTANT> To print Hello World in Python, use the print function:

<CODE>
print("Hello, World!")
<ENDCODE>

When you run this, it displays: Hello, World!

<USER> Write a function to add two numbers
<ASSISTANT> Here is a simple function to add two numbers:

<CODE>
def add(a, b):
    return a + b

result = add(5, 3)
print(result)
<ENDCODE>

This function takes two parameters and returns their sum.

<USER> Explain what a function is
<ASSISTANT> A function is a reusable block of code that performs a specific task. Functions help organize code and avoid repetition.

Here is the basic structure:

<CODE>
def function_name(parameters):
    # Code block
    return result
<ENDCODE>

Parts of a function:
- def: keyword to define a function
- function_name: what you call your function
- parameters: input values
- return: sends back a result

Example:

<CODE>
def greet(name):
    message = f"Hello, {name}!"
    return message

print(greet("Alice"))
<ENDCODE>

<USER> What is a list?
<ASSISTANT> A list is an ordered, mutable collection in Python that can hold items of any type.

Creating lists:

<CODE>
my_list = []
numbers = [1, 2, 3, 4, 5]
fruits = ["apple", "banana", "cherry"]
mixed = [1, "hello", 3.14, True]
<ENDCODE>

Common operations:

<CODE>
fruits = ["apple", "banana"]

fruits.append("cherry")
fruits.insert(0, "apricot")

print(fruits[0])
print(fruits[-1])

fruits.remove("banana")
fruits.pop()

print(len(fruits))
<ENDCODE>

<USER> Write a fibonacci function
<ASSISTANT> Here is a Fibonacci function:

<CODE>
def fibonacci(n):
    if n <= 0:
        return 0
    if n == 1:
        return 1
    return fibonacci(n-1) + fibonacci(n-2)

for i in range(10):
    print(fibonacci(i), end=" ")
<ENDCODE>

Iterative version:

<CODE>
def fibonacci_iterative(n):
    if n <= 0:
        return 0
    if n == 1:
        return 1
    
    a, b = 0, 1
    for _ in range(2, n + 1):
        a, b = b, a + b
    return b
<ENDCODE>

<USER> What is recursion?
<ASSISTANT> Recursion is when a function calls itself to solve a problem.

Key components:
1. Base case - stops the recursion
2. Recursive case - function calls itself

Example:

<CODE>
def countdown(n):
    if n <= 0:
        print("Done!")
        return
    print(n)
    countdown(n - 1)

countdown(5)
<ENDCODE>

Factorial example:

<CODE>
def factorial(n):
    if n <= 1:
        return 1
    return n * factorial(n - 1)

print(factorial(5))
<ENDCODE>

<USER> Write bubble sort
<ASSISTANT> Here is the Bubble Sort algorithm:

<CODE>
def bubble_sort(arr):
    n = len(arr)
    
    for i in range(n):
        swapped = False
        
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        
        if not swapped:
            break
    
    return arr

numbers = [64, 34, 25, 12, 22, 11, 90]
print("Sorted:", bubble_sort(numbers.copy()))
<ENDCODE>

<USER> Write binary search
<ASSISTANT> Here is Binary Search:

<CODE>
def binary_search(arr, target):
    left = 0
    right = len(arr) - 1
    
    while left <= right:
        mid = (left + right) // 2
        
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    
    return -1

numbers = [1, 3, 5, 7, 9, 11, 13, 15]
print(binary_search(numbers, 7))
<ENDCODE>

<USER> What is a class?
<ASSISTANT> A class is a blueprint for creating objects:

<CODE>
class Dog:
    def __init__(self, name, age):
        self.name = name
        self.age = age
    
    def bark(self):
        return f"{self.name} says Woof!"

my_dog = Dog("Buddy", 3)
print(my_dog.bark())
<ENDCODE>

<USER> How do I read a file?
<ASSISTANT> Here is how to read files:

<CODE>
with open("filename.txt", "r") as file:
    content = file.read()
    print(content)

with open("filename.txt", "r") as file:
    for line in file:
        print(line.strip())
<ENDCODE>

<USER> How do I handle errors?
<ASSISTANT> Use try-except blocks:

<CODE>
try:
    result = 10 / 0
except ZeroDivisionError:
    print("Cannot divide by zero!")
<ENDCODE>

Multiple exceptions:

<CODE>
try:
    number = int(input("Enter number: "))
    result = 10 / number
except ValueError:
    print("Not a valid number!")
except ZeroDivisionError:
    print("Cannot divide by zero!")
except Exception as e:
    print(f"Error: {e}")
<ENDCODE>

<USER> Thanks!
<ASSISTANT> You are welcome! Feel free to ask me anything anytime! Happy coding!

<USER> Hi
<ASSISTANT> Hello! How can I help you today? I can write code, explain concepts, or answer your programming questions!

<USER> What is the difference between list and tuple?
<ASSISTANT> The main difference is mutability:
- Lists are mutable (can change)
- Tuples are immutable (cannot change)

<CODE>
my_list = [1, 2, 3]
my_tuple = (1, 2, 3)

my_list[0] = 100
my_list.append(4)
<ENDCODE>

Use lists for changing data, tuples for fixed data.
'''


class VedaTrainer:
    """Trainer for Veda Programming Assistant"""
    
    def __init__(self, vocab_size: int = 8000, max_length: int = 512, batch_size: int = 4):
        self.vocab_size = vocab_size
        self.max_length = max_length
        self.batch_size = batch_size
        self.tokenizer = VedaTokenizer(vocab_size=vocab_size)
        self.model = None
    
    def prepare_data(self, extra_data: str = ""):
        """Prepare training data"""
        data = TRAINING_DATA
        if extra_data:
            data += "\n\n" + extra_data
        
        if os.path.exists("programming.txt"):
            with open("programming.txt", 'r', encoding='utf-8') as f:
                code_data = f.read()
            data += "\n\n" + code_data
        
        self.tokenizer.fit([data])
        
        all_tokens = self.tokenizer.encode(data)
        print(f"Total tokens: {len(all_tokens)}")
        
        sequences = []
        stride = self.max_length // 2
        
        for i in range(0, len(all_tokens) - self.max_length - 1, stride):
            seq = all_tokens[i:i + self.max_length + 1]
            if len(seq) == self.max_length + 1:
                sequences.append(seq)
        
        if len(sequences) < 10:
            stride = self.max_length // 4
            sequences = []
            for i in range(0, len(all_tokens) - self.max_length - 1, stride):
                seq = all_tokens[i:i + self.max_length + 1]
                if len(seq) == self.max_length + 1:
                    sequences.append(seq)
        
        print(f"Created {len(sequences)} training sequences")
        
        sequences = np.array(sequences)
        X = sequences[:, :-1]
        y = sequences[:, 1:]
        
        dataset = tf.data.Dataset.from_tensor_slices((X, y))
        dataset = dataset.shuffle(1000).batch(self.batch_size).prefetch(1)
        
        return dataset
    
    def build_model(self):
        """Build the model"""
        self.model = VedaProgrammingLLM(
            vocab_size=self.tokenizer.vocabulary_size,
            max_length=self.max_length,
            d_model=256,
            num_heads=8,
            num_layers=4,
            ff_dim=512
        )
        
        self.model.compile(
            optimizer=keras.optimizers.Adam(1e-4),
            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=['accuracy']
        )
        
        dummy = tf.zeros((1, self.max_length), dtype=tf.int32)
        self.model(dummy)
        
        return self.model
    
    def train(self, epochs: int = 15, save_path: str = None, extra_data: str = ""):
        """Train the model"""
        if save_path is None:
            save_path = MODEL_DIR
        
        dataset = self.prepare_data(extra_data)
        self.build_model()
        
        self.model.summary()
        
        os.makedirs(save_path, exist_ok=True)
        
        history = self.model.fit(dataset, epochs=epochs, verbose=1)
        
        self.model.save_weights(os.path.join(save_path, "weights.h5"))
        self.tokenizer.save(os.path.join(save_path, "tokenizer.json"))
        
        config = self.model.get_config()
        with open(os.path.join(save_path, "config.json"), 'w') as f:
            json.dump(config, f)
        
        print(f"Model saved to {save_path}")
        return history
    
    def generate_response(self, user_input: str, max_tokens: int = 200, temperature: float = 0.7) -> str:
        """Generate a response"""
        prompt = f"<USER> {user_input}\n<ASSISTANT>"
        
        tokens = self.tokenizer.encode(prompt)
        
        generated = self.model.generate(
            tokens,
            max_new_tokens=max_tokens,
            temperature=temperature,
            repetition_penalty=1.2
        )
        
        response = self.tokenizer.decode(generated)
        
        if "<ASSISTANT>" in response:
            response = response.split("<ASSISTANT>")[-1].strip()
        if "<USER>" in response:
            response = response.split("<USER>")[0].strip()
        
        return response


if __name__ == "__main__":
    trainer = VedaTrainer()
    trainer.train(epochs=20)
    
    print("\nTesting:")
    tests = ["Hello!", "What is a function?"]
    for test in tests:
        print(f"User: {test}")
        print(f"Assistant: {trainer.generate_response(test)}")