import os
import re
import random
import math
import yaml
from typing import List, Dict, Tuple, Set, Any
from collections import defaultdict, Counter
import pandas as pd
from tqdm import tqdm

from feather import FeatherManager, similarity_score, calculate_confidence_score

class GrammarRules:
    @staticmethod
    def apply_all_rules(text: str) -> str:
        if not text:
            return text
        
        return text.strip()

class PatternExtractor:
    
    def __init__(self):
        pass
    
    def extract_keywords(self, text: str) -> List[str]:
        if not text:
            return []
        
        full_text_normalized = re.sub(r'\s+', ' ', text.strip().lower())
        words = re.findall(r'\b[a-zA-Z]+\b', full_text_normalized)
        
        result = [full_text_normalized]
        result.extend(words)
        
        return list(set(result))
    
    def create_pattern(self, user_input: str) -> str:
        if not user_input:
            return ""
        
        pattern = re.sub(r'\s+', ' ', user_input.strip().lower())
        
        pattern = f" {pattern} "
        
        return pattern
    
    def calculate_pattern_similarity(self, pattern1: str, pattern2: str) -> float:
        return similarity_score(pattern1.strip(), pattern2.strip())


class MiniModelTrainer:
    
    def __init__(self, feather_manager: FeatherManager):
        self.feather_manager = feather_manager
        self.pattern_extractor = PatternExtractor()
        self.grammar_rules = GrammarRules()
    
    def train_mini_model(self, training_pairs: List[Tuple[str, str]], confidence_threshold: float = 0.1) -> Dict[str, Any]:
        if not training_pairs or len(training_pairs) < 2:
            return None

        keyword_patterns = []
        responses = []
        pattern_confidences = []
        all_keywords = []
        
        for user_input, ai_response in training_pairs:
            processed_response = ai_response.strip()
            
            # Get both the pattern and keywords
            pattern = self.pattern_extractor.create_pattern(user_input)
            keywords = self.pattern_extractor.extract_keywords(user_input)
            all_keywords.extend(keywords)
            
            # Add the main pattern
            keyword_patterns.append(pattern)
            responses.append(processed_response)
            individual_confidence = min(0.9, len(training_pairs) / 20.0)
            pattern_confidences.append(individual_confidence)
        
        if not keyword_patterns:
            return None
        
        base_confidence = min(0.9, len(training_pairs) / 20.0)
        keyword_counter = Counter(all_keywords)
        top_keywords = [word for word, count in keyword_counter.most_common(10)]
        
        mini_model = {
            'patterns': keyword_patterns,
            'responses': responses,
            'pattern_confidences': pattern_confidences,
            'confidence': base_confidence,
            'grammar_rules': [],
            'keywords': top_keywords,
            'training_samples': len(training_pairs)
        }
        
        return mini_model
    
    def should_merge_models(self, model1: Dict[str, Any], model2: Dict[str, Any], merge_threshold: float = 0.8) -> bool:
        keywords1 = set(model1.get('keywords', []))
        keywords2 = set(model2.get('keywords', []))
        
        if not keywords1 or not keywords2:
            return False
        
        keyword_similarity = len(keywords1.intersection(keywords2)) / len(keywords1.union(keywords2))
        
        responses1 = model1.get('responses', [])
        responses2 = model2.get('responses', [])
        
        response_similarities = []
        for r1 in responses1[:5]:
            for r2 in responses2[:5]:
                sim = similarity_score(r1, r2)
                response_similarities.append(sim)
        
        avg_response_similarity = sum(response_similarities) / len(response_similarities) if response_similarities else 0
        
        min_confidence = min(model1.get('confidence', 0), model2.get('confidence', 0))
        
        return (keyword_similarity > merge_threshold and 
                avg_response_similarity > merge_threshold and 
                min_confidence > 0.7)
    
    def merge_mini_models(self, model1: Dict[str, Any], model2: Dict[str, Any]) -> Dict[str, Any]:
        patterns1 = model1.get('patterns', [])
        patterns2 = model2.get('patterns', [])
        responses1 = model1.get('responses', [])
        responses2 = model2.get('responses', [])
        confidences1 = model1.get('pattern_confidences', [1.0] * len(patterns1))
        confidences2 = model2.get('pattern_confidences', [1.0] * len(patterns2))
        
        merged_model = {
            'patterns': patterns1 + patterns2,
            'responses': responses1 + responses2,
            'pattern_confidences': confidences1 + confidences2,
            'confidence': (model1.get('confidence', 0) + model2.get('confidence', 0)) / 2,
            'grammar_rules': list(set(model1.get('grammar_rules', []) + model2.get('grammar_rules', []))),
            'keywords': list(set(model1.get('keywords', []) + model2.get('keywords', []))),
            'training_samples': model1.get('training_samples', 0) + model2.get('training_samples', 0)
        }
        
        return merged_model


class AgGPTTrainer:
    
    def __init__(self, models_dir: str = "models"):
        self.feather_manager = FeatherManager(models_dir)
        self.mini_trainer = MiniModelTrainer(self.feather_manager)
        self.target_size_mb = 5
        self.estimated_size_per_pair = 1000
        self.chunk_size = (self.target_size_mb * 1024 * 1024) // self.estimated_size_per_pair
        self.readable_weights_dir = "readable_weights"
        os.makedirs(self.readable_weights_dir, exist_ok=True)
    
    def save_model_as_yaml(self, model_data: Dict[str, Any], model_id: int):
        try:
            filename = f"AgGPT_Model_{model_id:04d}.yaml"
            filepath = os.path.join(self.readable_weights_dir, filename)
            
            print(f"Creating YAML data for model {model_id}...")
            yaml_data = {
                'model_info': {
                    'model_id': model_id,
                    'confidence': model_data.get('confidence', 0.5),
                    'training_samples': model_data.get('training_samples', 0),
                    'keywords': model_data.get('keywords', [])
                },
                'patterns_and_responses': []
            }
            
            patterns = model_data.get('patterns', [])
            responses = model_data.get('responses', [])
            weights = model_data.get('weights', [])
            
            print(f"Processing {len(patterns)} patterns...")
            for i in range(len(patterns)):
                entry = {
                    'pattern': patterns[i] if i < len(patterns) else '',
                    'response': responses[i] if i < len(responses) else '',
                    'weight': weights[i] if i < len(weights) else 1.0
                }
                yaml_data['patterns_and_responses'].append(entry)
            
            print(f"Writing YAML to {filepath}...")
            with open(filepath, 'w', encoding='utf-8') as f:
                yaml.dump(yaml_data, f, default_flow_style=False, allow_unicode=True, indent=2)
            
            print(f"Saved readable model: {filename}")
        except Exception as e:
            print(f"Error in save_model_as_yaml: {e}")
            import traceback
            traceback.print_exc()
    
    def load_training_data(self, file_path: str) -> List[Tuple[str, str]]:
        training_pairs = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        conversations = content.split('<eos>')
        
        print(f"Processing {len(conversations)} conversation chunks...")
        
        for conversation in tqdm(conversations, desc="Parsing conversations"):
            conversation = conversation.strip()
            if not conversation:
                continue
            
            user_match = re.search(r'user:\s*(.*?)(?=\n<pad>|\nai:|$)', conversation, re.DOTALL)
            ai_match = re.search(r'ai:\s*(.*?)$', conversation, re.DOTALL)
            
            if user_match and ai_match:
                user_input = user_match.group(1).strip()
                ai_response = ai_match.group(1).strip()
                
                user_input = re.sub(r'<pad>', '', user_input).strip()
                ai_response = re.sub(r'<pad>', '', ai_response).strip()
                
                if user_input and ai_response and len(user_input) > 0 and len(ai_response) > 0:
                    training_pairs.append((user_input, ai_response))
        
        print(f"Extracted {len(training_pairs)} training pairs")
        return training_pairs
    
    def create_training_chunks(self, training_pairs: List[Tuple[str, str]]) -> List[List[Tuple[str, str]]]:
        shuffled_pairs = training_pairs.copy()
        random.shuffle(shuffled_pairs)
        
        chunks = []
        total_pairs = len(shuffled_pairs)
        
        for i in range(0, total_pairs, self.chunk_size):
            chunk = shuffled_pairs[i:i + self.chunk_size]
            if len(chunk) >= 5:
                chunks.append(chunk)
        
        print(f"Created {len(chunks)} training chunks (target: {self.target_size_mb}MB each)")
        return chunks
    
    def train_multiple_corpora(self, training_files: List[str] = None, merge_similar: bool = True):
        """Train on multiple corpora files sequentially"""
        if training_files is None:
            training_files = ["training_data/corpora.txt", "training_data/corpora2.txt"]
        
        print("Starting AgGPT-18 Multi-Corpora Training with Scalable Feather Architecture")
        print("=" * 70)
        
        cleared_count = self.feather_manager.clear_all_models()
        if cleared_count > 0:
            print(f"Cleared {cleared_count} existing models")
        
        all_trained_models = []
        total_model_id = 1
        
        for file_idx, training_file in enumerate(training_files, 1):
            print(f"\n--- Training on file {file_idx}/{len(training_files)}: {training_file} ---")
            
            if not os.path.exists(training_file):
                print(f"Warning: Training file {training_file} does not exist. Skipping...")
                continue
            
            if os.path.getsize(training_file) == 0:
                print(f"Warning: Training file {training_file} is empty. Skipping...")
                continue
            
            print(f"Loading training data from {training_file}...")
            training_pairs = self.load_training_data(training_file)
            
            if not training_pairs:
                print(f"No training data found in {training_file}. Skipping...")
                continue
            
            print(f"Creating training chunks for {training_file}...")
            training_chunks = self.create_training_chunks(training_pairs)
            
            print(f"Training mini-models from {training_file}...")
            file_trained_models = []
            
            progress_bar = tqdm(training_chunks, desc=f"Training from {os.path.basename(training_file)}")
            for chunk_idx, chunk in enumerate(progress_bar):
                print(f"\nProcessing chunk {chunk_idx + 1}/{len(training_chunks)}")
                mini_model = self.mini_trainer.train_mini_model(chunk)
                
                if mini_model:
                    file_trained_models.append(mini_model)
                    all_trained_models.append(mini_model)
                    print(f"Saving model {total_model_id}...")
                    self.feather_manager.save_mini_model(mini_model, total_model_id)
                    
                    if total_model_id == 1:
                        print("Saving first model as YAML...")
                        try:
                            self.save_model_as_yaml(mini_model, total_model_id)
                            print("YAML saved successfully")
                        except Exception as e:
                            print(f"Error saving YAML: {e}")
                    
                    total_model_id += 1
                    print(f"Model {total_model_id - 1} completed")
                    
                    try:
                        progress_bar.set_postfix({
                            'File Models': len(file_trained_models),
                            'Total Models': len(all_trained_models),
                            'Confidence': f"{mini_model['confidence']:.3f}"
                        })
                    except Exception as e:
                        print(f"Error updating progress bar: {e}")
            
            print(f"Completed training on {training_file}: {len(file_trained_models)} mini-models created")
            print(f"Total models so far: {len(all_trained_models)}")
            
            if merge_similar and len(all_trained_models) > 1:
                print(f"Merging similar models after processing {training_file}...")
                self._merge_similar_models()
                current_count = self.feather_manager.get_model_count()
                print(f"Models after merging: {current_count}")
        
        print(f"\n--- Multi-Corpora Training Complete ---")
        final_count = self.feather_manager.get_model_count()
        print(f"Final model count: {final_count}")
        print(f"Trained on {len([f for f in training_files if os.path.exists(f) and os.path.getsize(f) > 0])} corpora files")
        print("=" * 70)
    
    def train(self, training_file: str = "training_data/corpora.txt", merge_similar: bool = True):
        print("Starting AgGPT-18 Training with Scalable Feather Architecture")
        print("=" * 60)
        
        cleared_count = self.feather_manager.clear_all_models()
        if cleared_count > 0:
            print(f"Cleared {cleared_count} existing models")
        
        print("Loading training data...")
        training_pairs = self.load_training_data(training_file)
        
        if not training_pairs:
            print("No training data found!")
            return
        
        print("Creating training chunks...")
        training_chunks = self.create_training_chunks(training_pairs)
        
        print("Training mini-models...")
        trained_models = []
        model_id = 1
        
        progress_bar = tqdm(training_chunks, desc="Training mini-models")
        for chunk in progress_bar:
            mini_model = self.mini_trainer.train_mini_model(chunk)
            
            if mini_model:
                trained_models.append(mini_model)
                self.feather_manager.save_mini_model(mini_model, model_id)
                
                if model_id == 1:
                    self.save_model_as_yaml(mini_model, model_id)
                
                model_id += 1
                
                progress_bar.set_postfix({
                    'Models': len(trained_models),
                    'Confidence': f"{mini_model['confidence']:.3f}"
                })
        
        print(f"Trained {len(trained_models)} mini-models")
        
        if merge_similar and len(trained_models) > 1:
            print("Merging similar models...")
            self._merge_similar_models()
        
        final_count = self.feather_manager.get_model_count()
        print(f"Training complete! Final model count: {final_count}")
        print("=" * 60)
    
    def _merge_similar_models(self):
        all_models = self.feather_manager.load_all_models()
        if len(all_models) < 2:
            return
        
        merged_pairs = []
        models_to_delete = set()
        
        print(f"Checking {len(all_models)} models for merging opportunities...")
        
        progress_bar = tqdm(range(len(all_models)), desc="Merging models")
        for i in progress_bar:
            if i in models_to_delete:
                continue
            
            for j in range(i + 1, len(all_models)):
                if j in models_to_delete:
                    continue
                
                model1 = all_models[i]
                model2 = all_models[j]
                
                if self.mini_trainer.should_merge_models(model1, model2):
                    merged_model = self.mini_trainer.merge_mini_models(model1, model2)
                    
                    new_id = self.feather_manager.get_next_model_id()
                    self.feather_manager.save_mini_model(merged_model, new_id)
                    
                    models_to_delete.add(i)
                    models_to_delete.add(j)
                    merged_pairs.append((model1.get('model_id', i), model2.get('model_id', j), new_id))
                    
                    break
        
        for model_idx in models_to_delete:
            if model_idx < len(all_models):
                model_id = all_models[model_idx].get('model_id', model_idx + 1)
                self.feather_manager.delete_model(model_id)
        
        if merged_pairs:
            print(f"Merged {len(merged_pairs)} pairs of similar models")
        else:
            print("No similar models found for merging")


def main():
    print("AgGPT-18 Scalable Feather Architecture Trainer")
    print("=" * 50)
    
    trainer = AgGPTTrainer()
    
    try:
        trainer.train_multiple_corpora(merge_similar=True)
    except KeyboardInterrupt:
        print("\nTraining interrupted by user")
    except Exception as e:
        print(f"Training failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()