import streamlit as st
import numpy as np
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import Dataset
import time
from datetime import datetime
import plotly.graph_objects as go
from huggingface_hub import HfApi, HfFolder

# Initialize Hugging Face Authentication
def huggingface_login():
    token = st.text_input("Hugging Face Token", type="password")
    if token:
        HfFolder.save_token(token)
        api = HfApi()
        user_info = api.whoami(token)
        st.sidebar.write(f"Logged in as: {user_info['name']}")
        return token
    else:
        st.warning("Please enter your Hugging Face token")
        return None

# Advanced Cyberpunk Styling
def setup_advanced_cyberpunk_style():
    st.markdown("""
        <style>
        @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;500;700&display=swap');
        @import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&display=swap');
        .main-title { font-family: 'Orbitron', sans-serif; font-size: 40px; color: #00ffea; }
        /* Additional CSS styling for dashboard, progress bar, and background */
        </style>
    """, unsafe_allow_html=True)

# Initialize Model and Tokenizer
def initialize_model():
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

    # Set padding token to eos_token
    tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer

# Prepare Dataset
def prepare_dataset(data, tokenizer, block_size=128):
    def tokenize_function(examples):
        return tokenizer(examples['text'], truncation=True, max_length=block_size, padding='max_length')

    raw_dataset = Dataset.from_dict({'text': data})
    tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=['text'])
    tokenized_dataset = tokenized_dataset.map(lambda examples: {'labels': examples['input_ids']}, batched=True)
    tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
    return tokenized_dataset

# Training Dashboard Class
class TrainingDashboard:
    def __init__(self):
        self.metrics = {
            'current_loss': 0,
            'best_loss': float('inf'),
            'generation': 0,
            'start_time': time.time(),
            'training_speed': 0
        }
        self.history = []
        
    def update(self, loss, generation):
        self.metrics['current_loss'] = loss
        self.metrics['generation'] = generation
        if loss < self.metrics['best_loss']:
            self.metrics['best_loss'] = loss
        elapsed_time = time.time() - self.metrics['start_time']
        self.metrics['training_speed'] = generation / elapsed_time
        self.history.append({'loss': loss, 'timestamp': datetime.now().strftime('%H:%M:%S')})
    
    def display(self):
        st.write(f"**Generation:** {self.metrics['generation']}")
        st.write(f"**Current Loss:** {self.metrics['current_loss']:.4f}")
        st.write(f"**Best Loss:** {self.metrics['best_loss']:.4f}")
        st.write(f"**Training Speed:** {self.metrics['training_speed']:.2f} generations/sec")

# Display Progress Bar
def display_progress(progress):
    st.markdown(f"""
        <div class="progress-bar-container">
            <div class="progress-bar" style="width: {progress * 100}%"></div>
        </div>
    """, unsafe_allow_html=True)

# Custom Genetic Algorithm
class GeneticAlgorithm:
    def __init__(self, model, tokenizer, dataset, population_size, mutation_rate=0.1):
        self.model = model
        self.tokenizer = tokenizer
        self.dataset = dataset
        self.population_size = population_size
        self.mutation_rate = mutation_rate
        self.population = [self.clone_model() for _ in range(population_size)]
    
    def clone_model(self):
        # Create a clone of the model
        return GPT2LMHeadModel.from_pretrained("gpt2")

    def evaluate_fitness(self, model):
        # Calculate the loss for a given model on the dataset
        trainer = Trainer(
            model=model,
            args=TrainingArguments(output_dir="./results", per_device_train_batch_size=2, num_train_epochs=1),
            train_dataset=self.dataset,
            data_collator=DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False),
        )
        train_result = trainer.train()
        return train_result.training_loss

    def select_best_models(self, num_best=2):
        # Selects the top models based on fitness (loss)
        fitness_scores = [(self.evaluate_fitness(model), model) for model in self.population]
        fitness_scores.sort(key=lambda x: x[0])  # Sort by loss
        best_models = [model for _, model in fitness_scores[:num_best]]
        return best_models

    def crossover(self, parent1, parent2):
        # Perform crossover by combining layers from both parents
        child = self.clone_model()
        for (child_param, param1, param2) in zip(child.parameters(), parent1.parameters(), parent2.parameters()):
            # Randomly choose parameters from each parent based on crossover probability
            if np.random.rand() > 0.5:
                child_param.data = param1.data.clone()
            else:
                child_param.data = param2.data.clone()
        return child

    def mutate(self, model):
        # Mutate model by slightly adjusting its weights
        for param in model.parameters():
            if np.random.rand() < self.mutation_rate:
                mutation_tensor = torch.randn_like(param) * 0.02
                param.data += mutation_tensor

    def generate_new_population(self):
        best_models = self.select_best_models()
        new_population = []
        while len(new_population) < self.population_size:
            parent1, parent2 = np.random.choice(best_models, 2, replace=False)
            child = self.crossover(parent1, parent2)
            self.mutate(child)
            new_population.append(child)
        self.population = new_population

# Training Loop with Genetic Algorithm and Loading Screen
def training_loop(dashboard, ga, num_generations):
    with st.spinner("Training in progress..."):
        for generation in range(1, num_generations + 1):
            best_loss = min([ga.evaluate_fitness(model) for model in ga.population])
            dashboard.update(best_loss, generation)
            progress = generation / num_generations
            display_progress(progress)
            dashboard.display()
            ga.generate_new_population()
            time.sleep(0.5)  # Simulate delay for each generation

# Main Function
def main():
    setup_advanced_cyberpunk_style()
    st.markdown('<h1 class="main-title">Neural Evolution GPT-2 Training Hub</h1>', unsafe_allow_html=True)

    # Hugging Face Account Login
    token = huggingface_login()
    if token is None:
        return

    # Load Model and Tokenizer
    model, tokenizer = initialize_model()

    # Prepare Data
    data = ["Sample training text"] * 10  # Replace with real data
    train_dataset = prepare_dataset(data, tokenizer)

    # Initialize Dashboard
    dashboard = TrainingDashboard()
    
    # Sidebar Configuration
    st.sidebar.markdown("### Training Parameters")
    num_generations = st.sidebar.slider("Generations", 1, 50, 10)
    population_size = st.sidebar.slider("Population Size", 4, 20, 10)
    mutation_rate = st.sidebar.slider("Mutation Rate", 0.01, 0.5, 0.1)
    
    # Initialize Genetic Algorithm
    ga = GeneticAlgorithm(model, tokenizer, train_dataset, population_size, mutation_rate)
    
    # Run Training
    if st.button("Start Training"):
        training_loop(dashboard, ga, num_generations)

if __name__ == "__main__":
    main()