import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup, PreTrainedModel, AutoConfig, PretrainedConfig
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
from typing import Dict, List, Optional
import numpy as np
from tqdm import tqdm
import wandb
import logging
import json
from torch.amp import autocast, GradScaler

class ChainOfThoughtDataset(Dataset):
    def __init__(self, hf_dataset, tokenizer, max_length: int = 512, split="train", num_examples: Optional[int] = None):
        self.dataset = hf_dataset[split]
        if num_examples is not None:
            self.dataset = self.dataset.select(range(min(num_examples, len(self.dataset))))
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        item = self.dataset[idx]
        document = item['document']
        summary = item['summary']

        input_text = f"Summarize: {document[:1000]}... Thought process:"
        target_text = f"1. Understand task\n2. Identify topics\n3. Extract key info\n4. Organize\n5. Condense\n6. Draft\n7. Review\n\nSummary: {summary[:100]}..."

        full_text = input_text + target_text

        encoded = self.tokenizer.encode_plus(
            full_text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )

        input_ids = encoded["input_ids"].squeeze(0)
        attention_mask = encoded["attention_mask"].squeeze(0)

        target_start = len(self.tokenizer.encode(input_text, add_special_tokens=False))

        labels = input_ids.clone()
        labels[:target_start] = -100

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels
        }

class ThoughtModule(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int = 4):
        super().__init__()
        self.self_attention = nn.MultiheadAttention(hidden_size, num_heads)
        self.norm = nn.LayerNorm(hidden_size)
        self.feed_forward = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 2),
            nn.GELU(),
            nn.Linear(hidden_size * 2, hidden_size)
        )

    def forward(self, x):
        attended, _ = self.self_attention(x, x, x)
        x = self.norm(x + attended)
        return self.norm(x + self.feed_forward(x))

class AdvancedChainOfThoughtConfig(PretrainedConfig):
    model_type = "advanced_cot"
    
    def __init__(
        self,
        base_model_name="gpt2",
        num_thought_steps=3,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.base_model_name = base_model_name
        self.num_thought_steps = num_thought_steps


class AdvancedChainOfThoughtModel(PreTrainedModel):
    config_class = AdvancedChainOfThoughtConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        
        # Load the base model for causal language modeling
        self.base_model = AutoModelForCausalLM.from_pretrained(config.base_model_name)
        
        # Load the tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name, use_fast=True)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # Get hidden_size from the base model's configuration
        hidden_size = self.base_model.config.hidden_size
        
        # Thought modules based on the base model's hidden size
        self.thought_modules = nn.ModuleList([ThoughtModule(hidden_size) for _ in range(config.num_thought_steps)])

        # Output layer to project hidden states to vocabulary size
        self.output_layer = nn.Linear(hidden_size, self.base_model.config.vocab_size)

    def forward(self, input_ids, attention_mask=None, labels=None):
        # Use the base model's transformer directly
        base_outputs = self.base_model.transformer(input_ids, attention_mask=attention_mask)
        hidden_states = base_outputs.last_hidden_state  # Get last hidden states from the base model

        # Apply thought modules on top of the base model's hidden states
        for thought_module in self.thought_modules:
            hidden_states = thought_module(hidden_states)

        # Compute logits
        lm_logits = self.output_layer(hidden_states)

        # Compute loss if labels are provided
        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))

        return lm_logits, loss

    @torch.no_grad()
    def generate(self, input_ids, attention_mask=None, **kwargs):
        return self.base_model.generate(input_ids, attention_mask=attention_mask, **kwargs)

    def get_output_embeddings(self):
        return self.output_layer

    def set_output_embeddings(self, new_embeddings):
        self.output_layer = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
        return self.base_model.prepare_inputs_for_generation(input_ids, past=past, **kwargs)

    def _reorder_cache(self, past, beam_idx):
        return self.base_model._reorder_cache(past, beam_idx)


    
def train_advanced_cot_model(
    model: AdvancedChainOfThoughtModel,
    train_dataloader: DataLoader,
    epochs: int = 1,
    lr: float = 5e-5,
    warmup_steps: int = 100,
    gradient_accumulation_steps: int = 4,
    max_grad_norm: float = 1.0,
    logging_steps: int = 10
):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    total_steps = len(train_dataloader) * epochs // gradient_accumulation_steps
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)

    scaler = GradScaler()

    wandb.init(project="advanced-cot-training")
    
    global_step = 0
    for epoch in range(epochs):
        model.train()
        epoch_loss = 0
        
        with tqdm(total=len(train_dataloader), desc=f"Epoch {epoch + 1}/{epochs}") as pbar:
            for step, batch in enumerate(train_dataloader):
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)

                with autocast('cuda'):
                    _, loss = model(input_ids, attention_mask=attention_mask, labels=labels)
                    loss = loss / gradient_accumulation_steps

                scaler.scale(loss).backward()

                epoch_loss += loss.item()

                if (step + 1) % gradient_accumulation_steps == 0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
                    scaler.step(optimizer)
                    scaler.update()
                    optimizer.zero_grad(set_to_none=True)
                    scheduler.step() 
                    global_step += 1

                    if global_step % logging_steps == 0:
                        wandb.log({
                            "train_loss": epoch_loss / (step + 1),
                            "learning_rate": scheduler.get_last_lr()[0],
                            "global_step": global_step
                        })

                pbar.update(1)
                pbar.set_postfix({"loss": epoch_loss / (step + 1)})

                # Clear CUDA cache to free up memory
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

    wandb.finish()

def prepare_dataset(tokenizer, batch_size=2, max_length=512, train_examples=20):
    dataset = load_dataset("multi_news")
    train_dataset = ChainOfThoughtDataset(dataset, tokenizer, max_length=max_length, split="train", num_examples=train_examples)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    return train_dataloader

def main():
    try:
        base_model_name = "gpt2"
        config = AdvancedChainOfThoughtConfig(base_model_name=base_model_name, num_thought_steps=6)
        cot_model = AdvancedChainOfThoughtModel(config)

        train_dataloader = prepare_dataset(
            cot_model.tokenizer,
            batch_size=1,
            max_length=512,
        )

        train_advanced_cot_model(
            cot_model, 
            train_dataloader, 
            epochs=1,
            lr=5e-5, 
            gradient_accumulation_steps=900
        )

        # Save the model and configuration
        cot_model.save_pretrained("advanced_cot_multi_news_model")
        cot_model.tokenizer.save_pretrained("advanced_cot_multi_news_model")

        # Test loading and generation
        loaded_model = AdvancedChainOfThoughtModel.from_pretrained("advanced_cot_multi_news_model")
        loaded_tokenizer = AutoTokenizer.from_pretrained("advanced_cot_multi_news_model")

        test_document = """
        The United Nations Climate Change Conference, COP26, is set to take place in Glasgow, Scotland. 
        World leaders, experts, and activists will gather to discuss and negotiate global climate action. 
        The conference aims to accelerate progress towards the goals of the Paris Agreement and the UN Framework Convention on Climate Change.
        Key topics will include reducing greenhouse gas emissions, adapting to climate impacts, and financing climate action in developing countries.
        """

        prompt = f"Summarize: {test_document}\nThought process:"
        input_ids = loaded_tokenizer.encode(prompt, return_tensors="pt")
        
        output_ids = loaded_model.generate(
            input_ids,
            max_length=300,
            num_beams=2,
            no_repeat_ngram_size=2,
            top_k=50,
            top_p=0.95,
            temperature=0.7,
            do_sample=True,
        )

        response = loaded_tokenizer.decode(output_ids[0], skip_special_tokens=True)
        print(response)

    except Exception as e:
        print(f"An error occurred: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()


An error occurred: 'AdvancedChainOfThoughtConfig' object has no attribute 'hidden_size'
Traceback (most recent call last):
  File "/home/ramachandra/Desktop/llm_testing/new2.py", line 224, in main
    cot_model = AdvancedChainOfThoughtModel(config)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Desktop/llm_testing/new2.py", line 103, in __init__
    hidden_size = self.base_model.config.hidden_size
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/.local/lib/python3.12/site-packages/transformers/configuration_utils.py", line 264, in __getattribute__
    return super().__getattribute__(key)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'AdvancedChainOfThoughtConfig' object has no attribute 'hidden_size'